import Foundation
import Network
import Combine
import SwiftUI

class MeshNetworkManager: NSObject, ObservableObject {
    
    // MARK: - Published State
    @Published var connectedPeerUsers: [PeerUser] = []
    @Published var receivedMessages: [MeshMessage] = []
    @Published var liveTranscripts: [String: String] = [:] // Keyed by senderNodeID
    
    // MARK: - Internal Types
    
    private struct TransportMessage: Codable {
        enum MessageType: Codable {
            // Added instance to identification for ghost filtering
            case identification(nodeID: String, instance: Int, name: String, role: UserRole, color: String)
            case broadcast(payload: Data)
        }
        let type: MessageType
        let senderNodeID: String
    }
    
    // MARK: - Properties
    
    private let serviceType = "_access-mesh._tcp" // Must match Constants.serviceType (implied)
    private let myNodeID: String
    private var myInstance: Int
    
    // UI Compatibility Accessors
    var myNodeIDPublic: String { myNodeID }
    var myInstancePublic: Int { myInstance }
    var isAtCapacity: Bool { connectedPeerUsers.count >= 7 }
    
    // User Identity
    private var myUserName: String = UIDevice.current.name
    var userName: String { myUserName }
    private var myUserColor: String = "#3357FF"
    private var myUserRole: UserRole = .hearing
    
    // Network Objects
    private var listener: NWListener?
    private var browser: NWBrowser?
    
    // Connection State
    // nodeID -> NWConnection
    private var connections: [String: NWConnection] = [:]
    
    // Track instances for ghost filtering
    private var peerInstances: [String: Int] = [:]
    
    // Stable peers (Traffic Gated)
    private var stableNodeIDs: Set<String> = []
    
    // Pending Connections (to avoid duplicates)
    private var pendingConnections: Set<String> = []
    
    // Internal Peer Map (before they become "connectedPeerUsers" or while unstable)
    // nodeID -> PeerUser
    private var internalPeerMap: [String: PeerUser] = [:]
    
    private var isRunning = false
    private let log = Logger()
    
    // MARK: - Initialization
    
    override init() {
        self.myNodeID = NodeIdentity.nodeID
        self.myInstance = NodeIdentity.nextInstance()
        super.init()
    }
    
    // MARK: - Public API
    
    func setIdentity(name: String, colorHex: String, role: UserRole) {
        self.myUserName = name
        self.myUserColor = colorHex
        self.myUserRole = role
    }
    
    func start() {
        guard !isRunning else { return }
        isRunning = true
        
        // Ensure fresh instance on every start
        self.myInstance = NodeIdentity.nextInstance()
        
        startListener()
        startBrowser()
        
        log.info("MeshNetworkManager started. NodeID: \(self.myNodeID), Instance: \(self.myInstance)")
    }
    
    func stop() {
        isRunning = false
        
        listener?.cancel()
        listener = nil
        
        browser?.cancel()
        browser = nil
        
        for (_, conn) in connections {
            conn.cancel()
        }
        connections.removeAll()
        pendingConnections.removeAll()
        stableNodeIDs.removeAll()
        internalPeerMap.removeAll()
        peerInstances.removeAll()
        
        DispatchQueue.main.async {
            self.connectedPeerUsers.removeAll()
            self.receivedMessages.removeAll()
            self.liveTranscripts.removeAll()
        }
        
        log.info("MeshNetworkManager stopped.")
    }
    
    func send(message: MeshMessage) {
        // Encode the app-level MeshMessage
        guard let payload = try? JSONEncoder().encode(message) else { return }
        
        // Wrap in TransportMessage
        let transportMsg = TransportMessage(type: .broadcast(payload: payload), senderNodeID: myNodeID)
        guard let data = try? JSONEncoder().encode(transportMsg) else { return }
        
        // Broadcast to all STABLE connections
        for (nodeID, connection) in connections where stableNodeIDs.contains(nodeID) {
            NetworkFraming.send(data, over: connection) { [weak self] error in
                if let error = error {
                    self?.log.error("Send failed to \(nodeID): \(error)")
                }
            }
        }
        
        // Handle local display of the message
        DispatchQueue.main.async {
            if message.isPartial {
                // Don't show own partials usually, or handle in UI
            } else if !message.isHandshake && !message.isKeepAlive {
                self.receivedMessages.append(message)
            }
        }
    }
    
    // MARK: - Network Parameters
    
    private func createParameters() -> NWParameters {
        let parameters = NWParameters.tcp
        parameters.includePeerToPeer = true // Enables AWDL
        
        // Enable KeepAlive to detect broken links
        let tcpOptions = NWProtocolTCP.Options()
        tcpOptions.enableKeepalive = true
        tcpOptions.keepaliveInterval = 10
        parameters.defaultProtocolStack.transportProtocol = tcpOptions
        
        return parameters
    }
    
    // MARK: - Listener (Host)
    
    private func startListener() {
        do {
            let parameters = createParameters()
            listener = try NWListener(using: parameters)
            
            // Advertise NodeID and Instance in TXT record
            let txtRecord = "nodeID=\(myNodeID)\tinstance=\(myInstance)"
            var txtData = Data()
            for entry in txtRecord.split(separator: "\t") {
                let s = String(entry)
                guard let d = s.data(using: .utf8), d.count < 255 else { continue }
                txtData.append(UInt8(d.count))
                txtData.append(d)
            }
            
            listener?.service = NWListener.Service(
                name: myNodeID, // Service Name = NodeID
                type: serviceType,
                txtRecord: txtData
            )
            
            listener?.newConnectionHandler = { [weak self] connection in
                self?.handleIncomingConnection(connection)
            }
            
            listener?.stateUpdateHandler = { [weak self] state in
                self?.log.info("Listener State: \(state)")
            }
            
            listener?.start(queue: .main)
            
        } catch {
            log.error("Failed to create listener: \(error)")
        }
    }
    
    private func handleIncomingConnection(_ connection: NWConnection) {
        connection.start(queue: .main)
        
        // Start receiving immediately to get the Identification message
        NetworkFraming.startReceiving(from: connection, onData: { [weak self] data in
            self?.handleData(data, from: connection)
        }, onError: { [weak self] error in
            self?.log.error("Receive error on incoming connection: \(error)")
            connection.cancel()
        })
    }
    
    // MARK: - Browser (Client)
    
    private func startBrowser() {
        let parameters = createParameters()
        browser = NWBrowser(for: .bonjour(type: serviceType, domain: nil), using: parameters)
        
        browser?.browseResultsChangedHandler = { [weak self] results, changes in
            self?.handleBrowseResults(results, changes: changes)
        }
        
        browser?.start(queue: .main)
    }
    
    private func handleBrowseResults(_ results: Set<NWBrowser.Result>, changes: Set<NWBrowser.Result.Change>) {
        for change in changes {
            switch change {
            case .added(let result):
                handleDiscovered(result)
            case .removed(let result):
                handleLost(result)
            default: break
            }
        }
    }
    
    private func handleDiscovered(_ result: NWBrowser.Result) {
        guard case .service(let name, _, _, _) = result.endpoint else { return }
        
        var peerNodeID = name // Default to service name
        var peerInstance = 0
        
        if case .bonjour(let txt) = result.metadata {
            if let nid = txt["nodeID"] { peerNodeID = nid }
            if let instStr = txt["instance"], let inst = Int(instStr) { peerInstance = inst }
        }
        
        guard peerNodeID != myNodeID else { return }
        
        // GHOST FILTER: Check if we know a newer instance
        if let existingInst = peerInstances[peerNodeID] {
             if peerInstance < existingInst {
                log.info("Ignoring older instance for \(peerNodeID): \(peerInstance) < \(existingInst)")
                return
             } else if peerInstance > existingInst {
                 log.info("Discovered NEWER instance \(peerInstance) > \(existingInst) for \(peerNodeID). Disconnecting stale connection.")
                 // Force disconnect old connection to allow re-initiation
                 if let conn = connections[peerNodeID] {
                     conn.cancel()
                     connections.removeValue(forKey: peerNodeID)
                     pendingConnections.remove(peerNodeID)
                     stableNodeIDs.remove(peerNodeID)
                     internalPeerMap.removeValue(forKey: peerNodeID)
                 }
             }
        }
        
        log.info("Discovered \(peerNodeID) (inst: \(peerInstance))")
        
        // Tie-Breaker: Lower ID initiates connection
        if myNodeID < peerNodeID {
            initiateConnection(to: result.endpoint, peerNodeID: peerNodeID)
        } else {
            log.info("Waiting for \(peerNodeID) to connect (I have higher ID)")
        }
    }
    
    private func handleLost(_ result: NWBrowser.Result) {
        if case .service(let name, _, _, _) = result.endpoint {
            log.info("Browser lost service: \(name)")
        }
    }
    
    // MARK: - Connection Management
    
    private func initiateConnection(to endpoint: NWEndpoint, peerNodeID: String) {
        guard connections[peerNodeID] == nil else { return }
        guard !pendingConnections.contains(peerNodeID) else { return }
        
        pendingConnections.insert(peerNodeID)
        log.info("Initiating connection to \(peerNodeID)")
        
        let parameters = createParameters()
        let connection = NWConnection(to: endpoint, using: parameters)
        
        connection.stateUpdateHandler = { [weak self] state in
            self?.handleConnectionState(state, peerNodeID: peerNodeID, connection: connection)
        }
        
        connection.start(queue: .main)
        connections[peerNodeID] = connection
        
        NetworkFraming.startReceiving(from: connection, onData: { [weak self] data in
            self?.handleData(data, from: connection)
        }, onError: { [weak self] error in
            self?.handleConnectionError(peerNodeID: peerNodeID, error: error)
        })
    }
    
    private func handleConnectionState(_ state: NWConnection.State, peerNodeID: String, connection: NWConnection) {
        switch state {
        case .ready:
            log.info("Connection ready: \(peerNodeID). Sending Identity.")
            sendIdentity(to: connection)
            pendingConnections.remove(peerNodeID)
        case .failed(let error):
            log.error("Connection failed \(peerNodeID): \(error)")
            handleConnectionError(peerNodeID: peerNodeID, error: error)
        case .cancelled:
            handleConnectionError(peerNodeID: peerNodeID, error: nil)
        default: break
        }
    }
    
    private func handleConnectionError(peerNodeID: String, error: Error?) {
        connections[peerNodeID]?.cancel()
        connections.removeValue(forKey: peerNodeID)
        pendingConnections.remove(peerNodeID)
        stableNodeIDs.remove(peerNodeID)
        internalPeerMap.removeValue(forKey: peerNodeID)
        
        updatePublicPeerList()
    }
    
    // MARK: - Data Handling & Logic
    
    private func sendIdentity(to connection: NWConnection) {
        let msg = TransportMessage(
            type: .identification(nodeID: myNodeID, instance: myInstance, name: myUserName, role: myUserRole, color: myUserColor),
            senderNodeID: myNodeID
        )
        guard let data = try? JSONEncoder().encode(msg) else { return }
        
        NetworkFraming.send(data, over: connection) { [weak self] error in
            if let error = error {
                self?.log.error("Failed to send identity: \(error)")
            }
        }
    }
    
    private func handleData(_ data: Data, from connection: NWConnection) {
        guard let msg = try? JSONDecoder().decode(TransportMessage.self, from: data) else { return }
        
        switch msg.type {
        case .identification(let nodeID, let instance, let name, let role, let color):
            handleIdentification(nodeID: nodeID, instance: instance, name: name, role: role, color: color, connection: connection)
            
        case .broadcast(let payload):
            handleBroadcast(payload, from: msg.senderNodeID)
        }
    }
    
    private func handleIdentification(nodeID: String, instance: Int, name: String, role: UserRole, color: String, connection: NWConnection) {
        log.info("Identified peer: \(nodeID) (inst: \(instance)) - \(name)")
        
        // GHOST CHECK
        if let existingInst = peerInstances[nodeID] {
            if instance < existingInst {
                log.info("Rejecting identification from OLDER instance \(instance) < \(existingInst)")
                connection.cancel()
                return
            } else if instance > existingInst {
                log.info("Upgrading to NEWER instance \(instance) > \(existingInst)")
                // Old connection should be cleaned up by overwrite
            }
        }
        
        peerInstances[nodeID] = instance
        
        // Store Connection if this was an incoming one (where we didn't know the ID yet)
        if connections[nodeID] == nil {
            connections[nodeID] = connection
            // We accepted incoming, so we must send our identity back if we haven't
            sendIdentity(to: connection)
        } else if connections[nodeID] !== connection {
            // Race condition: We have two connections for same NodeID.
            // Tie-breaker logic should prevent this, but if it happens:
            // "Lower ID initiates".
            // If I am Lower, I should have initiated. The one I initiated should stay.
            // If I am Higher, they should have initiated. This one (incoming) should stay.
            
            // However, usually we overwrite with the newest authenticated one.
            connections[nodeID]?.cancel() // Cancel old one
            connections[nodeID] = connection
        }
        
        // Create PeerUser
        let peer = PeerUser(nodeID: nodeID, name: name, colorHex: color, role: role, isStable: false)
        internalPeerMap[nodeID] = peer
        
        // Traffic Gate / Stabilization
        Task { @MainActor in
            let delay: UInt64 = 4_000_000_000
            try? await Task.sleep(nanoseconds: delay)
            
            // Verify connection still exists and matches this nodeID
            guard connections[nodeID] != nil else { return }
            
            stableNodeIDs.insert(nodeID)
            if var p = internalPeerMap[nodeID] {
                p.isStable = true
                internalPeerMap[nodeID] = p
            }
            
            updatePublicPeerList()
            log.info("Peer stabilized: \(name)")
        }
    }
    
    private func handleBroadcast(_ payload: Data, from senderNodeID: String) {
        guard let message = try? JSONDecoder().decode(MeshMessage.self, from: payload) else { return }
        
        DispatchQueue.main.async { [weak self] in
            guard let self = self else { return }
            
            if message.isPartial {
                if message.content.isEmpty {
                    self.liveTranscripts.removeValue(forKey: senderNodeID)
                } else {
                    self.liveTranscripts[senderNodeID] = message.content
                }
            } else if message.isKeepAlive {
                // Ignore
            } else if message.isHandshake {
                // Handshake logic is handled by TransportMessage.identification now
            } else {
                self.receivedMessages.append(message)
                self.liveTranscripts.removeValue(forKey: senderNodeID)
            }
        }
    }
    
    private func updatePublicPeerList() {
        DispatchQueue.main.async {
            self.connectedPeerUsers = self.internalPeerMap
                .filter { self.stableNodeIDs.contains($0.key) }
                .map { $0.value }
        }
    }
    
    // MARK: - Logger Helper
    struct Logger {
        func info(_ msg: String) { print("🕸️ [Mesh] \(msg)") }
        func error(_ msg: String) { print("🕸️ ❌ [Mesh] \(msg)") }
    }
}