import Foundation

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

/// RFC 5389 STUN binding, the minimum needed to learn this socket's server-reflexive
/// address (public ip:port as the NAT maps it). Pure encode/decode here — pinned by
/// `StunTests` vectors — with a small query driver over `UDPSocket` below.
///
/// Only Binding Request/Success and (XOR-)MAPPED-ADDRESS are implemented: no
/// authentication, no TURN — the servers queried are public, stateless echo services, and
/// the answer reveals nothing the network path doesn't already see.
public enum StunMessage {
    public static let magicCookie: UInt32 = 0x2112_A442
    public static let headerLength = 20
    static let bindingRequestType: UInt16 = 0x0001
    static let bindingSuccessType: UInt16 = 0x0101
    static let attrMappedAddress: UInt16 = 0x0001
    static let attrXORMappedAddress: UInt16 = 0x0020

    /// A 96-bit transaction ID — fresh random per query, echoed by the server; the response
    /// matcher drops anything that doesn't echo ours (off-path junk, stale retransmits).
    public static func makeTransactionID() -> Data {
        Data((0..<12).map { _ in UInt8.random(in: .min ... .max) })
    }

    /// A Binding Request: 20-byte header, zero attributes.
    public static func bindingRequest(transactionID: Data) -> Data {
        precondition(transactionID.count == 12, "STUN transaction ID must be 12 bytes")
        var out = Data(capacity: headerLength)
        out.append(contentsOf: withUnsafeBytes(of: bindingRequestType.bigEndian, Array.init))
        out.append(contentsOf: [0, 0]) // message length: no attributes
        out.append(contentsOf: withUnsafeBytes(of: magicCookie.bigEndian, Array.init))
        out.append(transactionID)
        return out
    }

    /// The mapped (public) address from a Binding Success, or nil for anything else —
    /// wrong type, wrong transaction, malformed attributes. XOR-MAPPED-ADDRESS wins over
    /// the legacy MAPPED-ADDRESS when both appear (RFC 5389 §12.2: XOR exists because
    /// ALGs rewrite literal addresses in payloads).
    public static func parseMappedAddress(_ response: Data, transactionID: Data) -> UDPEndpoint? {
        let bytes = [UInt8](response)
        guard bytes.count >= headerLength else { return nil }
        let type = UInt16(bytes[0]) << 8 | UInt16(bytes[1])
        guard type == bindingSuccessType else { return nil }
        let length = Int(UInt16(bytes[2]) << 8 | UInt16(bytes[3]))
        let cookie = UInt32(bytes[4]) << 24 | UInt32(bytes[5]) << 16 | UInt32(bytes[6]) << 8 | UInt32(bytes[7])
        guard cookie == magicCookie, Data(bytes[8..<20]) == transactionID,
              bytes.count >= headerLength + length else { return nil }

        var mapped: UDPEndpoint?
        var offset = headerLength
        let end = headerLength + length
        while offset + 4 <= end {
            let attrType = UInt16(bytes[offset]) << 8 | UInt16(bytes[offset + 1])
            let attrLen = Int(UInt16(bytes[offset + 2]) << 8 | UInt16(bytes[offset + 3]))
            let valueStart = offset + 4
            guard valueStart + attrLen <= end else { return mapped }
            let value = Array(bytes[valueStart..<valueStart + attrLen])
            switch attrType {
            case attrXORMappedAddress:
                if let endpoint = decodeAddress(value, xored: true, transactionID: transactionID) {
                    return endpoint // authoritative — stop scanning
                }
            case attrMappedAddress:
                if mapped == nil {
                    mapped = decodeAddress(value, xored: false, transactionID: transactionID)
                }
            default:
                break
            }
            // Attributes are 32-bit aligned; length excludes padding.
            offset = valueStart + attrLen + (4 - attrLen % 4) % 4
        }
        return mapped
    }

    private static func decodeAddress(_ value: [UInt8], xored: Bool, transactionID: Data) -> UDPEndpoint? {
        guard value.count >= 8 else { return nil }
        let family = value[1]
        var port = UInt16(value[2]) << 8 | UInt16(value[3])
        if xored { port ^= UInt16(magicCookie >> 16) }

        if family == 0x01 { // IPv4
            var octets = Array(value[4..<8])
            if xored {
                let cookie = withUnsafeBytes(of: magicCookie.bigEndian, Array.init)
                for i in 0..<4 { octets[i] ^= cookie[i] }
            }
            return .v4(ip: octets.map(String.init).joined(separator: "."), port: port)
        }
        if family == 0x02, value.count >= 20 { // IPv6: XOR with cookie ++ transaction ID
            var octets = Array(value[4..<20])
            if xored {
                let key = withUnsafeBytes(of: magicCookie.bigEndian, Array.init) + [UInt8](transactionID)
                for i in 0..<16 { octets[i] ^= key[i] }
            }
            var addr = in6_addr()
            withUnsafeMutableBytes(of: &addr) { $0.copyBytes(from: octets) }
            guard let ip = withUnsafePointer(to: &addr, { DirectSock.ntop(family: AF_INET6, $0) })
            else { return nil }
            return .v6(ip: ip, port: port)
        }
        return nil
    }
}

/// What two STUN mappings from different servers say about the NAT in front of this socket.
public enum NATClass: String, Sendable, Codable {
    /// Endpoint-independent mapping: both servers saw the same public port — hole punching
    /// works (full-cone / restricted variants all punch with simultaneous send).
    case cone
    /// The mapping changed per destination: punching at a learned port is mostly futile.
    /// Still *attempted* when only one side is symmetric — the cone side's port is stable.
    case symmetric
    /// One or zero servers answered — unclassified, punch optimistically.
    case unknown
}

/// The STUN query driver: resolve → send → match, with the RFC 5389 doubling retransmit
/// (500 ms initial RTO). `discover` runs the full two-server classification used by
/// candidate gathering.
public enum StunClient {
    public static let defaultServers = ["stun.cloudflare.com:3478", "stun.l.google.com:19302"]
    public static let serversEnvKey = "NUCLEIC_STUN_SERVERS"
    public static let serversDefaultsKey = "nucleic.sync.stunServers"

    /// The configured server list: env first (scripted runs / NAT sim), then the defaults
    /// key, then the built-in pair — same precedence as every other sync knob.
    public static func servers(
        environment: [String: String] = ProcessInfo.processInfo.environment,
        defaults: UserDefaults = .standard
    ) -> [String] {
        let raw = environment[serversEnvKey] ?? defaults.string(forKey: serversDefaultsKey)
        let parsed = raw?.split(separator: ",").map { $0.trimmingCharacters(in: .whitespaces) }
            .filter { !$0.isEmpty } ?? []
        return parsed.isEmpty ? defaultServers : parsed
    }

    public struct Discovery: Sendable {
        /// The reflexive endpoint per answering server, in query order.
        public let mappings: [UDPEndpoint]
        public let natClass: NATClass

        public var reflexive: UDPEndpoint? { mappings.first }
    }

    /// Query every server from the same socket and classify the NAT from the answers:
    /// same public port everywhere ⇒ cone (punchable), differing ⇒ symmetric. Servers
    /// are queried serially — the second query must observe a *second* outbound
    /// destination through the same local port, which is exactly what distinguishes the
    /// two NAT classes.
    public static func discover(
        socket: UDPSocket,
        servers serverList: [String]? = nil,
        timeoutPerAttempt: Duration = .milliseconds(500),
        attempts: Int = 3
    ) async -> Discovery {
        var mappings: [UDPEndpoint] = []
        for server in serverList ?? servers() {
            guard let endpoint = await resolve(server: server, ipv6: socket.isIPv6) else { continue }
            if let mapped = await query(
                socket: socket, server: endpoint,
                timeoutPerAttempt: timeoutPerAttempt, attempts: attempts)
            {
                mappings.append(mapped)
            }
        }
        let natClass: NATClass
        if mappings.count >= 2 {
            natClass = Set(mappings.map(\.port)).count == 1 ? .cone : .symmetric
        } else {
            natClass = .unknown
        }
        return Discovery(mappings: mappings, natClass: natClass)
    }

    /// One Binding round-trip with doubling retransmit; nil when every attempt times out.
    public static func query(
        socket: UDPSocket,
        server: UDPEndpoint,
        timeoutPerAttempt: Duration = .milliseconds(500),
        attempts: Int = 3
    ) async -> UDPEndpoint? {
        let clock = ContinuousClock()
        var timeout = timeoutPerAttempt
        for _ in 0..<max(1, attempts) {
            let txid = StunMessage.makeTransactionID()
            let request = StunMessage.bindingRequest(transactionID: txid)
            do {
                try socket.send(request, to: server)
            } catch {
                return nil
            }
            let deadline = clock.now + timeout
            // Drain until deadline: unrelated datagrams (punch probes, a previous
            // attempt's late answer) can interleave with the response we want.
            while let packet = try? await socket.receive(deadline: deadline, clock: clock) {
                if let mapped = StunMessage.parseMappedAddress(packet.data, transactionID: txid) {
                    return mapped
                }
            }
            timeout *= 2
        }
        return nil
    }

    /// `host:port` → a dialable endpoint (first A/AAAA record matching the socket family).
    /// getaddrinfo blocks, so it runs off the concurrency pool.
    static func resolve(server: String, ipv6: Bool) async -> UDPEndpoint? {
        guard let colon = server.lastIndex(of: ":"),
              let port = UInt16(server[server.index(after: colon)...])
        else { return nil }
        let host = String(server[..<colon])

        return await Task.detached(priority: .utility) { () -> UDPEndpoint? in
            var hints = addrinfo()
            hints.ai_family = ipv6 ? AF_INET6 : AF_INET
            #if canImport(Darwin) || os(Windows)
            hints.ai_socktype = SOCK_DGRAM
            #elseif canImport(Musl)
            hints.ai_socktype = Int32(SOCK_DGRAM)
            #else
            hints.ai_socktype = Int32(SOCK_DGRAM.rawValue)
            #endif
            var result: UnsafeMutablePointer<addrinfo>?
            guard getaddrinfo(host, String(port), &hints, &result) == 0, let info = result else {
                return nil
            }
            defer { freeaddrinfo(result) }
            var node: UnsafeMutablePointer<addrinfo>? = info
            while let current = node {
                if let sa = current.pointee.ai_addr {
                    var storage = sockaddr_storage()
                    withUnsafeMutableBytes(of: &storage) { dest in
                        dest.copyBytes(
                            from: UnsafeRawBufferPointer(
                                start: sa, count: Int(current.pointee.ai_addrlen)))
                    }
                    if let endpoint = UDPSocket.endpoint(from: &storage) { return endpoint }
                }
                node = current.pointee.ai_next
            }
            return nil
        }.value
    }
}
