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

    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 }
        connection.send(content: data, completion: .contentProcessed { [weak self] error in
            if error != nil { self?.stop() }
        })
    }

    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 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) }
        }
    }
}