import Foundation
import Network
import UltraMeshCore

final class PeerSession {
    let connection: NWConnection
    let outbound: Bool
    var remote: NodeRecord?
    var onMessage: ((PeerSession, WireMessage) -> Void)?
    var onStop: ((PeerSession) -> Void)?
    private let queue: DispatchQueue
    private var stopped = false
    private struct Pending { let data: Data; let bulk: Bool }
    private var priorityQueue: [Pending] = []
    private var bulkQueue: [Pending] = []
    private var bulkBytes = 0
    private var sending = false
    private static let maximumPendingBulkBytes = 512 * 1024

    init(connection: NWConnection, outbound: Bool, queue: DispatchQueue) {
        self.connection = connection; self.outbound = outbound; self.queue = queue
    }

    func start(local: NodeRecord) {
        connection.stateUpdateHandler = { [weak self] state in
            guard let self else { return }
            switch state {
            case .ready:
                self.send(.hello(local)); self.receiveMessage()
            case .failed, .cancelled: self.stop()
            default: break
            }
        }
        connection.start(queue: queue)
    }

    func send(_ message: WireMessage) {
        guard !stopped, let data = try? FrameCodec.encode(WireEnvelope(message: message)) else { return }
        let bulk: Bool
        if case .packet(let packet) = message { bulk = packet.trafficClass == .nativeIPv6 } else { bulk = false }
        if bulk {
            guard bulkBytes + data.count <= Self.maximumPendingBulkBytes else { return }
            bulkBytes += data.count; bulkQueue.append(.init(data: data, bulk: true))
        } else { priorityQueue.append(.init(data: data, bulk: false)) }
        pumpWrites()
    }

    func stop() {
        guard !stopped else { return }; stopped = true
        connection.cancel(); onStop?(self)
    }

    private func receiveMessage() {
        receiveExactly(4, accumulated: Data()) { [weak self] prefix in
            guard let self, let prefix,
                  let size = try? FrameCodec.bodyLength(from: prefix) else { self?.stop(); return }
            self.receiveExactly(size, accumulated: Data()) { [weak self] body in
                guard let self, let body,
                      let envelope = try? FrameCodec.decode(WireEnvelope.self, body: body),
                      envelope.version == WireEnvelope.currentVersion else { self?.stop(); return }
                self.onMessage?(self, envelope.message)
                self.receiveMessage()
            }
        }
    }

    private func pumpWrites() {
        guard !stopped, !sending else { return }
        let pending: Pending
        if !priorityQueue.isEmpty { pending = priorityQueue.removeFirst() }
        else if !bulkQueue.isEmpty { pending = bulkQueue.removeFirst(); bulkBytes -= pending.data.count }
        else { return }
        sending = true
        connection.send(content: pending.data, completion: .contentProcessed { [weak self] error in
            guard let self else { return }
            self.queue.async {
                self.sending = false
                if error != nil { self.stop() } else { self.pumpWrites() }
            }
        })
    }

    private func receiveExactly(_ count: Int, accumulated: Data, completion: @escaping (Data?) -> Void) {
        if accumulated.count == count { completion(accumulated); return }
        connection.receive(minimumIncompleteLength: 1, maximumLength: count - accumulated.count) { [weak self] data, _, complete, error in
            guard self != nil, error == nil, let data, !data.isEmpty else { completion(nil); return }
            let joined = accumulated + data
            if joined.count == count { completion(joined) }
            else if complete { completion(nil) }
            else { self?.receiveExactly(count, accumulated: joined, completion: completion) }
        }
    }
}