SwiftChain/SwiftChainCore/Models/Node.swift

//
//  Node.swift
//  SwiftChain
//
//  Created by Victor on 27/11/2024.
//

import Foundation
import Network

public class Node {
    // Network properties
    private var peers: [NWConnection] = []
    private let port: NWEndpoint.Port = 8333
    private var listener: NWListener?
    private let queue = DispatchQueue(label: "com.blockchain.node")

    // Blockchain components
    private let blockchain: Blockchain
    private let accountManager: AccountManager
    private var pendingTransactions: [Transaction] = []
    private let maxTransactionsPerBlock: Int = 10

    public init() {
        self.accountManager = AccountManager()
        self.blockchain = Blockchain()
        setupListener()
    }

    enum MessageType: String, Codable {
        case transaction = "TX"
        case blockchainRequest = "BC_REQ"
        case blockchainResponse = "BC_RES"
        case newBlock = "NEW_BLOCK"
    }

    struct NetworkMessage: Codable {
        let type: MessageType
        let data: Data
    }

    // Configuration du listener
    private func setupListener() {
        let parameters = NWParameters.tcp
        do {
            listener = try NWListener(using: parameters, on: port)
            listener?.stateUpdateHandler = { [weak self] state in
                switch state {
                case .ready:
                    print("Node is listening on port \(self?.port.rawValue ?? 0)")
                case .failed(let error):
                    print("Listener failed with error: \(error)")
                default:
                    break
                }
            }

            listener?.newConnectionHandler = { [weak self] connection in
                self?.handleNewConnection(connection)
                print("New incoming connection from: \(connection.endpoint)")
            }

            listener?.start(queue: queue)
        } catch {
            print("Failed to create listener: \(error)")
        }
    }

    public func mineBlock(minerAddress: String) -> Block? {
        guard let lastBlock = blockchain.chain.last else {
            print("No existing blockchain")
            return nil
        }

        // Utiliser la méthode locale au lieu de memPool
        var transactions = getTransactionsForBlock()
        let miningReward = Transaction(
            sender: "SYSTEM",
            receiver: minerAddress,
            amount: 50,
            type: "MINING_REWARD"
        )
        transactions.append(miningReward)

        let newBlock = Block(
            transactions: transactions,
            previousHash: lastBlock.hash,
            index: blockchain.chain.count,
            difficulty: blockchain.calculateNewDifficulty()
        )

        newBlock.miner = minerAddress
        newBlock.timestamp = Int(Date().timeIntervalSince1970)
        let miningTime = newBlock.mineBlock()

        if accountManager.processBlock(newBlock) {
            blockchain.chain.append(newBlock)
            broadcastBlock(newBlock)

            // Nettoyer le mempool après avoir ajouté le bloc
            cleanupMempool(chain: blockchain.chain)

            print("""
            Block \(newBlock.index) created:
            Hash: \(newBlock.hash)
            Previous hash: \(newBlock.previousHash)
            Transactions: \(newBlock.transactions.count)
            Mining time: \(String(format: "%.2f", miningTime)) seconds
            """)
            return newBlock
        }

        return nil
    }

    // Public interface
    public func submitTransaction(_ transaction: Transaction) -> Bool {
        if addTransaction(transaction) {
            broadcastTransaction(transaction)
            return true
        }
        return false
    }

    public func getBalance(_ address: String) -> Int {
        return accountManager.getBalance(address)
    }

    public func getPendingTransactions() -> [Transaction] {
        return pendingTransactions
    }

    public func isChainValid() -> Bool {
        return blockchain.validateChain(blockchain.chain)
    }

    // Connexion à un pair
    public func connectToPeer(host: String) {
        let endpoint = NWEndpoint.hostPort(host: NWEndpoint.Host(host), port: port)
        let connection = NWConnection(to: endpoint, using: .tcp)

        connection.stateUpdateHandler = { [weak self] state in
            switch state {
            case .ready:
                print("Connected to peer: \(host)")
                self?.peers.append(connection)
                self?.startReceiving(connection)
                // Demander la blockchain au pair
                self?.requestBlockchain(from: connection)
            case .failed(let error):
                print("Connection failed: \(error)")
            case .cancelled:
                print("Connection cancelled: \(host)")
            default:
                break
            }
        }

        connection.start(queue: queue)
    }

    // Méthode pour lister les pairs connectés
    public func listPeers() {
        if peers.isEmpty {
            print("Aucun pair connecté")
            return
        }

        print("\nPairs connectés (\(peers.count)):")
        for (index, peer) in peers.enumerated() {
            print("\(index + 1). \(peer.endpoint)")
        }
    }

    // Propagation d'une transaction
    func broadcastTransaction(_ transaction: Transaction) {
        do {
            let transactionData = try JSONEncoder().encode(transaction)
            let message = NetworkMessage(type: .transaction, data: transactionData)
            let messageData = try JSONEncoder().encode(message)

            print("Broadcasting transaction to \(peers.count) peers")

            for peer in peers {
                peer.send(content: messageData, completion: .contentProcessed { error in
                    if let error = error {
                        print("Failed to send transaction: \(error)")
                    } else {
                        print("Transaction sent successfully to: \(peer.endpoint)")
                    }
                })
            }
        } catch {
            print("Failed to encode transaction: \(error)")
        }
    }

    // Réception des messages
    private func startReceiving(_ connection: NWConnection) {
        connection.receive(minimumIncompleteLength: 1, maximumLength: 65536) { [weak self] content, _, isComplete, error in
            if let data = content {
                self?.handleReceivedData(data, from: connection)
            }

            if let error = error {
                print("Receive error: \(error)")
            }

            if !isComplete, error == nil {
                self?.startReceiving(connection)
            }
        }
    }

    private func handleReceivedData(_ data: Data, from connection: NWConnection) {
        do {
            print("Received data of size: \(data.count) bytes")
            let message = try JSONDecoder().decode(NetworkMessage.self, from: data)
            print("Message decoded successfully, type: \(message.type)")

            switch message.type {
            case .transaction:
                let transaction = try JSONDecoder().decode(Transaction.self, from: message.data)
                print("Transaction decoded: \(transaction.sender) -> \(transaction.receiver): \(transaction.amount)")
                if addTransaction(transaction) { // Utiliser la méthode locale
                    print("Transaction added to mempool successfully")
                } else {
                    print("Failed to add transaction to mempool")
                }

            case .blockchainRequest:
                // Envoyer notre blockchain au pair
                sendBlockchain(to: connection)

            case .blockchainResponse:
                // Recevoir et traiter la blockchain
                let receivedChain = try JSONDecoder().decode([Block].self, from: message.data)
                handleReceivedBlockchain(receivedChain)

            case .newBlock:
                // Recevoir et traiter un nouveau bloc
                let block = try JSONDecoder().decode(Block.self, from: message.data)
                handleReceivedBlock(block)
            }
        } catch {
            print("Failed to decode received data: \(error)")
        }
    }

    private func handleNewConnection(_ connection: NWConnection) {
        connection.start(queue: queue)
        peers.append(connection)
        startReceiving(connection)
        print("New connection established with: \(connection.endpoint)")

        // Envoyer notre blockchain au nouveau pair
        sendBlockchain(to: connection)
    }

    private func requestBlockchain(from peer: NWConnection) {
        do {
            let message = NetworkMessage(type: .blockchainRequest, data: Data())
            let messageData = try JSONEncoder().encode(message)

            peer.send(content: messageData, completion: .contentProcessed { error in
                if let error = error {
                    print("Failed to request blockchain: \(error)")
                } else {
                    print("Blockchain request sent to: \(peer.endpoint)")
                }
            })
        } catch {
            print("Failed to encode blockchain request: \(error)")
        }
    }

    private func sendBlockchain(to peer: NWConnection) {
        do {
            let chainData = try JSONEncoder().encode(blockchain.chain)
            let message = NetworkMessage(type: .blockchainResponse, data: chainData)
            let messageData = try JSONEncoder().encode(message)

            peer.send(content: messageData, completion: .contentProcessed { error in
                if let error = error {
                    print("Failed to send blockchain: \(error)")
                } else {
                    print("Blockchain sent to: \(peer.endpoint)")
                }
            })
        } catch {
            print("Failed to encode blockchain: \(error)")
        }
    }

    private func handleReceivedBlockchain(_ receivedChain: [Block]) {
        print("Received blockchain with \(receivedChain.count) blocks")

        if receivedChain.count > blockchain.chain.count {
            print("Received chain is longer, validating...")
            if blockchain.validateChain(receivedChain) {
                // Sauvegarder l'état actuel
                let oldChain = blockchain.chain

                // Essayer d'appliquer la nouvelle chaîne
                blockchain.chain = receivedChain

                // Vérifier que toutes les transactions sont valides
                var isValid = true
                for block in receivedChain where !accountManager.processBlock(block) {
                    isValid = false
                    break
                }

                if isValid {
                    print("Blockchain updated successfully")
                    cleanupMempool(chain: receivedChain)
                } else {
                    blockchain.chain = oldChain
                    print("Failed to process transactions in received chain")
                }
            } else {
                print("Received chain is invalid")
            }
        } else {
            print("Current chain is longer or equal, keeping it")
        }
    }

    private func handleReceivedBlock(_ block: Block) {
        print("Received new block: \(block.hash)")

        // Vérifier que le bloc suit bien notre dernier bloc
        guard let lastBlock = blockchain.chain.last else {
            print("No existing chain")
            return
        }

        if block.previousHash != lastBlock.hash {
            print("Block does not connect to our chain")
            return
        }

        if !block.isValid() {
            print("Block is not valid")
            return
        }

        // Traiter les transactions du bloc
        if !accountManager.processBlock(block) {
            print("Failed to process block transactions")
            return
        }

        // Nettoyer le mempool des transactions incluses dans le bloc
        cleanupMempool(chain: [block])

        // Ajouter le bloc à la chaîne
        blockchain.chain.append(block)
        print("New block added to chain")

        // Propager le bloc aux autres pairs (sauf celui qui nous l'a envoyé)
        broadcastBlock(block)
    }

    private func broadcastBlock(_ block: Block) {
        do {
            let blockData = try JSONEncoder().encode(block)
            let message = NetworkMessage(type: .newBlock, data: blockData)
            let messageData = try JSONEncoder().encode(message)

            print("Broadcasting block to peers")

            for peer in peers {
                peer.send(content: messageData, completion: .contentProcessed { error in
                    if let error = error {
                        print("Failed to send block: \(error)")
                    } else {
                        print("Block sent successfully to: \(peer.endpoint)")
                    }
                })
            }
        } catch {
            print("Failed to encode block: \(error)")
        }
    }
}

// MARK: - Mempool methods

extension Node {
    // MemPool functionality
    private func addTransaction(_ transaction: Transaction) -> Bool {
        if validateTransaction(transaction) {
            pendingTransactions.append(transaction)
            print("Transaction added to mempool successfully")
            return true
        }
        return false
    }

    private func getTransactionsForBlock() -> [Transaction] {
        var validTransactions: [Transaction] = []
        var remainingTransactions: [Transaction] = []

        for transaction in pendingTransactions {
            if validTransactions.count >= maxTransactionsPerBlock {
                break
            }

            if validateTransaction(transaction) {
                validTransactions.append(transaction)
            } else {
                remainingTransactions.append(transaction)
            }
        }

        // Update pending transactions
        pendingTransactions = remainingTransactions + pendingTransactions.dropFirst(validTransactions.count)

        return validTransactions
    }

    private func cleanupMempool(chain: [Block]) {
        var remainingTransactions: [Transaction] = []

        for pendingTx in pendingTransactions {
            // On vérifie si la transaction est dans un des blocs
            var isInChain = false

            for block in chain {
                if block.transactions.contains(where: { chainTx in
                    chainTx.sender == pendingTx.sender &&
                        chainTx.receiver == pendingTx.receiver &&
                        chainTx.amount == pendingTx.amount &&
                        chainTx.type == pendingTx.type
                }) {
                    isInChain = true
                    break
                }
            }

            if !isInChain {
                remainingTransactions.append(pendingTx)
            }
        }

        print("Mempool cleanup: removed \(pendingTransactions.count - remainingTransactions.count) transactions")
        pendingTransactions = remainingTransactions
    }

    var hasPendingTransactions: Bool {
        return !pendingTransactions.isEmpty
    }

    public func printMemPoolStatus() {
        print("""
        MemPool Status:
        - Pending transactions: \(pendingTransactions.count)
        - Maximum transactions per block: \(maxTransactionsPerBlock)
        """)
    }

    private func checkDoubleSpending(_ transaction: Transaction) -> Bool {
        let address = transaction.sender
        let pendingAmount = pendingTransactions
            .filter { $0.sender == address }
            .reduce(0) { $0 + $1.amount }

        let currentBalance = accountManager.getBalance(address)

        if currentBalance < (pendingAmount + transaction.amount) {
            print("MemPool: Transaction refusée - tentative de double dépense détectée")
            return false
        }

        return true
    }

    private func validateTransaction(_ transaction: Transaction) -> Bool {
        if !transaction.isValid() {
            return false
        }

        if transaction.type == "MINING_REWARD" {
            return true
        }

        if !checkDoubleSpending(transaction) {
            return false
        }

        if !accountManager.canProcessTransaction(transaction) {
            print("MemPool: Transaction refusée - solde insuffisant")
            return false
        }

        print("MemPool: Transaction validée avec succès")
        return true
    }
}