import Foundation

#if canImport(Darwin)
import Darwin
#elseif canImport(Musl)
import Musl
#elseif canImport(Glibc)
import Glibc
#elseif canImport(WinSDK)
import WinSDK
#endif

import NucleicProtocol

/// `nucleic-punch-harness` — drives the real STUN + hole-punch core (Covalence direct,
/// SYNC_PROTOCOL §3.3) so `scripts/nat-sim` can exercise it across simulated NATs (Linux
/// network namespaces + nftables masquerade). It is deliberately the *production* code path:
/// `StunClient`, `DirectCandidateGatherer`, and `PunchEngine` are what run here, so a green
/// sim is evidence about the shipping punch logic, not a re-implementation.
///
/// Three sub-commands, wired by the sim's `run.sh`:
///
///   stun-server            Reflect the observed source as XOR-MAPPED-ADDRESS (a minimal
///                          RFC 5389 responder) so a peer namespace can classify its NAT
///                          without reaching the public internet.
///     --port N
///
///   classify               Gather candidates through the configured STUN servers and print
///                          the NAT class + reflexive mapping as one JSON line. Proves the
///                          symmetric-vs-cone detection the punch policy keys off.
///     (reads NUCLEIC_STUN_SERVERS)
///
///   punch --role a|b --peer-file PATH --token HEX --peer-token HEX
///                          One side of a punch. Writes its candidates to `--out`, waits for
///                          the peer's `--peer-file`, then punches. Exits 0 iff a round trip
///                          validated within the budget — the sim asserts on the exit code.
@main
struct PunchHarness {
    static func main() async {
        let args = Array(CommandLine.arguments.dropFirst())
        guard let command = args.first else {
            fail("usage: nucleic-punch-harness <stun-server|classify|punch> [options]")
        }
        switch command {
        case "stun-server": await runStunServer(Options(args))
        case "classify": await runClassify(Options(args))
        case "punch": await runPunch(Options(args))
        default: fail("unknown command \(command)")
        }
    }

    // MARK: stun-server

    static func runStunServer(_ options: Options) async {
        let port = UInt16(options["port"] ?? "3478") ?? 3478
        let bind = options["bind"]
        guard let socket = try? UDPSocket(port: port, ipv6: false, bindAddress: bind) else {
            fail("stun-server: could not bind \(bind ?? "*"):\(port)")
        }
        note("stun-server listening on \(bind ?? "*"):\(socket.localPort)")
        let clock = ContinuousClock()
        while true {
            guard let packet = try? await socket.receive(deadline: clock.now + .seconds(3600))
            else { continue }
            guard packet.data.count >= StunMessage.headerLength else { continue }
            let txid = Data(packet.data.suffix(12))
            let response = successResponse(source: packet.from, transactionID: txid)
            try? socket.send(response, to: packet.from)
        }
    }

    /// A Binding Success carrying XOR-MAPPED-ADDRESS = the observed source (v4 only — the
    /// sim's NATs are v4). Mirrors what a public STUN server returns.
    static func successResponse(source: UDPEndpoint, transactionID: Data) -> Data {
        var out = Data([0x01, 0x01, 0x00, 0x0C, 0x21, 0x12, 0xA4, 0x42])
        out += transactionID
        out += Data([0x00, 0x20, 0x00, 0x08])
        let xport = source.port ^ UInt16(StunMessage.magicCookie >> 16)
        out += Data([0x00, 0x01, UInt8(xport >> 8), UInt8(xport & 0xFF)])
        let octets = source.ip.split(separator: ".").compactMap { UInt8($0) }
        let cookie: [UInt8] = [0x21, 0x12, 0xA4, 0x42]
        out += Data(zip(octets, cookie).map(^))
        return out
    }

    // MARK: classify

    static func runClassify(_ options: Options) async {
        guard let sockets = try? DirectCandidateGatherer.openSockets() else {
            fail("classify: could not open sockets")
        }
        defer { sockets.close() }
        let (candidates, natClass) = await DirectCandidateGatherer.gather(sockets: sockets)
        let reflexive = candidates.first { $0.kind == .reflexive }
        let json = """
        {"natClass":"\(natClass.rawValue)","reflexive":"\(reflexive.map { "\($0.ip):\($0.port)" } ?? "none")","candidates":\(candidates.count)}
        """
        print(json)
        exit(natClass == .unknown ? 2 : 0)
    }

    // MARK: punch

    static func runPunch(_ options: Options) async {
        guard let role = options["role"], let out = options["out"],
              let peerFile = options["peer-file"],
              let tokenHex = options["token"], let localToken = Data(hex: tokenHex),
              let peerTokenHex = options["peer-token"], let remoteToken = Data(hex: peerTokenHex)
        else {
            fail("punch: need --role --out --peer-file --token --peer-token")
        }
        let budget = Duration.seconds(Int64(options["budget"].flatMap(Int.init) ?? 12))

        guard let sockets = try? DirectCandidateGatherer.openSockets() else {
            fail("punch[\(role)]: could not open sockets")
        }
        let (candidates, natClass) = await DirectCandidateGatherer.gather(sockets: sockets)
        note("punch[\(role)] natClass=\(natClass.rawValue) localPort=\(sockets.v4.localPort) "
            + "candidates=[\(candidates.map { "\($0.kind.rawValue):\($0.ip):\($0.port)" }.joined(separator: ","))]")

        // Publish our candidates for the peer, then wait for theirs. A file drop is the
        // sim's shared channel (Noise isn't needed at this layer — the sim isolates the
        // punch mechanics).
        writeCandidates(candidates, to: out)
        guard let peerCandidates = await waitForCandidates(peerFile, timeout: .seconds(15)) else {
            fail("punch[\(role)]: peer candidates never arrived")
        }
        note("punch[\(role)] punching peer=[\(peerCandidates.map { "\($0.kind.rawValue):\($0.ip):\($0.port)" }.joined(separator: ","))]")

        guard punchWorthAttempting(local: natClass, remote: .unknown) else {
            note("punch[\(role)]: both-symmetric skip")
            sockets.close()
            exit(3)
        }

        let engine = PunchEngine(
            sockets: sockets, localToken: localToken, remoteToken: remoteToken,
            remoteCandidates: peerCandidates)
        var validated: UDPEndpoint?
        for await update in await engine.run(budget: budget) where update.roundTrip {
            validated = update.endpoint
            break
        }
        sockets.close()
        if let validated {
            note("punch[\(role)]: VALIDATED \(validated.description)")
            exit(0)
        }
        note("punch[\(role)]: no round trip within budget")
        exit(1)
    }

    // MARK: candidate file exchange

    static func writeCandidates(_ candidates: [DirectCandidate], to path: String) {
        let lines = candidates.map { "\($0.kind.rawValue) \($0.ip) \($0.port)" }.joined(separator: "\n")
        try? (lines + "\n").write(toFile: path, atomically: true, encoding: .utf8)
    }

    static func waitForCandidates(_ path: String, timeout: Duration) async -> [DirectCandidate]? {
        let clock = ContinuousClock()
        let deadline = clock.now + timeout
        while clock.now < deadline {
            if let text = try? String(contentsOfFile: path, encoding: .utf8), !text.isEmpty {
                let candidates = text.split(separator: "\n").compactMap { line -> DirectCandidate? in
                    let parts = line.split(separator: " ")
                    guard parts.count == 3, let kind = DirectCandidate.Kind(rawValue: String(parts[0])),
                          let port = UInt16(parts[2]) else { return nil }
                    return DirectCandidate(kind: kind, ip: String(parts[1]), port: port)
                }
                if !candidates.isEmpty { return candidates }
            }
            try? await clock.sleep(for: .milliseconds(100))
        }
        return nil
    }

    // MARK: helpers

    static func note(_ message: String) {
        FileHandle.standardError.write(Data((message + "\n").utf8))
    }

    static func fail(_ message: String) -> Never {
        note("error: " + message)
        exit(64)
    }
}

/// Bare `--key value` option parsing (the sim passes explicit flags).
struct Options {
    private var values: [String: String] = [:]

    init(_ args: [String]) {
        var i = 0
        while i < args.count {
            let arg = args[i]
            if arg.hasPrefix("--") {
                let key = String(arg.dropFirst(2))
                if i + 1 < args.count, !args[i + 1].hasPrefix("--") {
                    values[key] = args[i + 1]
                    i += 2
                } else {
                    values[key] = ""
                    i += 1
                }
            } else {
                i += 1
            }
        }
    }

    subscript(_ key: String) -> String? { values[key] }
}

extension Data {
    init?(hex: String) {
        guard hex.count % 2 == 0 else { return nil }
        var bytes = [UInt8]()
        var index = hex.startIndex
        while index < hex.endIndex {
            let next = hex.index(index, offsetBy: 2)
            guard let byte = UInt8(hex[index..<next], radix: 16) else { return nil }
            bytes.append(byte)
            index = next
        }
        self.init(bytes)
    }
}
