import Foundation

/// Temporary opt-in tracing for the direct-upgrade handshake (set NUCLEIC_DIRECT_DEBUG=1).
enum DirectDebug {
    static let on = ProcessInfo.processInfo.environment["NUCLEIC_DIRECT_DEBUG"] == "1"
    static func log(_ s: String) {
        if on { FileHandle.standardError.write(Data("DIRECT: \(s)\n".utf8)) }
    }
}

/// The Covalence direct-path kill switch (SYNC_PROTOCOL §3.3): on by default; `NUCLEIC_DIRECT
/// =off` (env) or `nucleic.sync.direct = "off"` (defaults) forces relay-only — for debugging,
/// or networks where UDP probing is unwelcome. Mirrors `SyncTransportSetting`'s
/// env-then-defaults precedence.
public enum DirectPathSetting {
    public static let envKey = "NUCLEIC_DIRECT"
    public static let defaultsKey = "nucleic.sync.direct"

    public static func enabled(
        environment: [String: String] = ProcessInfo.processInfo.environment,
        defaults: UserDefaults = .standard
    ) -> Bool {
        let raw = environment[envKey] ?? defaults.string(forKey: defaultsKey)
        switch raw?.lowercased() {
        case "off", "0", "false", "no": return false
        default: return true
        }
    }

    /// Whether this build can carry the direct data path. The RUDP transport is pure POSIX
    /// UDP, so — unlike the earlier QUIC cut — every platform qualifies (Mac, iOS, and Linux
    /// `nucleicd` alike).
    public static var platformSupported: Bool { true }
}

/// One side's live punch attempt: the sockets, the gathered candidates, and the token we
/// advertised (which the peer echoes in punches and stamps on RUDP datagrams).
public struct DirectAttemptState: Sendable {
    public let sockets: DirectCandidateGatherer.Sockets
    public let candidates: [DirectCandidate]
    public let natClass: NATClass
    public let localToken: Data

    /// Open sockets, run STUN, gather everything one `directOffer`/`directAnswer` needs.
    public static func gather(stunServers: [String]? = nil) async -> DirectAttemptState? {
        guard let sockets = try? DirectCandidateGatherer.openSockets() else { return nil }
        let (candidates, natClass) = await DirectCandidateGatherer.gather(
            sockets: sockets, stunServers: stunServers)
        guard !candidates.isEmpty else {
            sockets.close()
            return nil
        }
        return DirectAttemptState(
            sockets: sockets, candidates: candidates, natClass: natClass,
            localToken: PunchWire.makeToken())
    }

    public var offer: DirectPunchOffer {
        DirectPunchOffer(candidates: candidates, natClass: natClass, token: localToken)
    }

    /// The socket for a given validated remote's family (its RUDP send socket).
    public func socket(for remote: UDPEndpoint) -> UDPSocket {
        remote.isIPv6 ? (sockets.v6 ?? sockets.v4) : sockets.v4
    }
}

/// Host side of one direct upgrade (SYNC_PROTOCOL §3.3), owned per relay connection. Replies
/// route back through the connection's sealed sender (the closures). Lifecycle:
/// `start(remoteOffer:)` → punch → `select(_:)` → RUDP over the live punched socket → `directGo`
/// → attach + verify + switch on the mux. No TLS, no port handoff: the punched socket keeps
/// carrying bytes, now as RUDP, and Noise stays the end-to-end security boundary above it.
public actor DirectHostUpgrader {
    private let mux: MultipathFrameChannel
    private let sendAnswer: @Sendable (DirectPunchOffer) async -> Void
    private let sendGo: @Sendable () async -> Void
    private let sendDecline: @Sendable (String) async -> Void

    private var state: DirectAttemptState?
    private var clientToken: Data?
    /// The client source we validated, keyed by address family (v6 → true). The client's
    /// `select` names which of *our* candidates it chose, so we must reply on the same family
    /// — a machine can validate both a v4 and a v6 pair, and the two ends must agree on one.
    private var validatedByFamily: [Bool: UDPEndpoint] = [:]
    private var engine: PunchEngine?
    private var engineTask: Task<Void, Never>?
    private var finished = false

    public init(
        mux: MultipathFrameChannel,
        sendAnswer: @escaping @Sendable (DirectPunchOffer) async -> Void,
        sendGo: @escaping @Sendable () async -> Void,
        sendDecline: @escaping @Sendable (String) async -> Void
    ) {
        self.mux = mux
        self.sendAnswer = sendAnswer
        self.sendGo = sendGo
        self.sendDecline = sendDecline
    }

    public func start(remoteOffer: DirectPunchOffer) async {
        guard !finished, state == nil else {
            await sendDecline("direct attempt already in flight")
            return
        }
        guard let gathered = await DirectAttemptState.gather() else {
            await sendDecline("no direct candidates on this host")
            return
        }
        guard punchWorthAttempting(local: gathered.natClass, remote: remoteOffer.natClass) else {
            gathered.sockets.close()
            await sendDecline("both ends behind symmetric NAT")
            return
        }
        state = gathered
        clientToken = remoteOffer.token
        await sendAnswer(gathered.offer)

        // Punch toward the client's candidates. The first working remote (round-trip OR just
        // inbound — either proves a reachable return address) becomes our RUDP peer.
        let engine = PunchEngine(
            sockets: gathered.sockets, localToken: gathered.localToken,
            remoteToken: remoteOffer.token, remoteCandidates: remoteOffer.candidates)
        self.engine = engine
        engineTask = Task { [weak self] in
            for await update in await engine.run(budget: .seconds(20)) {
                await self?.noteValidated(update.endpoint)
            }
        }
    }

    private func noteValidated(_ endpoint: UDPEndpoint) {
        if validatedByFamily[endpoint.isIPv6] == nil { validatedByFamily[endpoint.isIPv6] = endpoint }
    }

    public func select(_ selected: DirectSelect) async {
        guard !finished, let state, let clientToken else {
            await sendDecline("no direct attempt to select")
            return
        }
        // Reply on the family the client selected: it validated one of our candidates, and its
        // RUDP will run on that family's socket, so ours must too.
        let selectedIsV6 = selected.ip.contains(":")
        let fallback: UDPEndpoint = selectedIsV6
            ? .v6(ip: selected.ip, port: selected.port) : .v4(ip: selected.ip, port: selected.port)
        let remote = validatedByFamily[selectedIsV6]
            ?? validatedByFamily[!selectedIsV6] ?? fallback
        DirectDebug.log("host.select selectedIsV6=\(selectedIsV6) validated=\(validatedByFamily.values.map(\.description)) using=\(remote.description)")
        // Stop punching before RUDP adopts the socket — one reader per socket. `stop()` is
        // what actually halts the engine's receive loop (cancelling `engineTask` alone
        // wouldn't; see `PunchEngine.stop`).
        await engine?.stop()
        engineTask?.cancel()
        engine = nil
        engineTask = nil

        let rudp = RUDPFrameChannel(
            socket: state.socket(for: remote), remote: remote,
            sendToken: clientToken, recvToken: state.localToken)
        mux.attachDirect(rudp)
        await sendGo()
        let switched = await mux.verifyAndSwitch()
        DirectDebug.log("host.select verifyAndSwitch → \(switched)")
        await finish(switched ? nil : "direct path failed verification")
    }

    private func finish(_ failure: String?) async {
        guard !finished else { return }
        finished = true
        await engine?.stop()
        engineTask?.cancel()
        engine = nil
        engineTask = nil
        if let failure {
            state?.sockets.close()
            await sendDecline(failure)
        }
        // On success the socket is owned by the RUDP channel (attached to the mux); leave it.
        state = nil
    }

    public func cancel() async {
        await finish(nil)
    }
}

/// Client side of one direct upgrade, owned per relay session. Lifecycle: `makeOffer()`
/// (send as `directOffer`) → `answer(_:)` (host's candidates: punch until a round trip) →
/// returns the winner to send as `directSelect` → `go()` (RUDP over the punched socket,
/// attach + verify + switch).
public actor DirectClientUpgrader {
    private let mux: MultipathFrameChannel
    private var state: DirectAttemptState?
    private var hostToken: Data?
    private var engine: PunchEngine?
    private var winner: UDPEndpoint?
    private var finished = false

    public init(mux: MultipathFrameChannel) {
        self.mux = mux
    }

    public func makeOffer() async -> DirectPunchOffer? {
        guard !finished, state == nil, mux.canAttachDirect else { return nil }
        guard let gathered = await DirectAttemptState.gather() else { return nil }
        state = gathered
        return gathered.offer
    }

    /// Punch the host's candidates until a round trip proves one. Returns the candidate to
    /// `directSelect` (nil ⇒ the attempt failed; caller backs off on relay).
    public func answer(_ hostOffer: DirectPunchOffer, budget: Duration = .seconds(12)) async -> DirectSelect? {
        guard !finished, let state, winner == nil else { return nil }
        guard punchWorthAttempting(local: state.natClass, remote: hostOffer.natClass) else {
            await abandon()
            return nil
        }
        hostToken = hostOffer.token

        let engine = PunchEngine(
            sockets: state.sockets, localToken: state.localToken,
            remoteToken: hostOffer.token, remoteCandidates: hostOffer.candidates)
        self.engine = engine
        var validated: UDPEndpoint?
        for await update in await engine.run(budget: budget) where update.roundTrip {
            validated = update.endpoint
            break
        }
        guard let validated else {
            await abandon()
            return nil
        }
        winner = validated
        return DirectSelect(endpoint: validated)
    }

    /// The host is ready: build RUDP over the punched socket to the winner, attach + verify +
    /// switch. Returns whether the session now rides direct.
    public func go() async -> Bool {
        guard !finished, let state, let winner, let hostToken else { return false }
        // Halt the punch engine before RUDP adopts the socket (see `PunchEngine.stop`).
        await engine?.stop()
        engine = nil
        DirectDebug.log("client.go winner=\(winner.description)")
        let rudp = RUDPFrameChannel(
            socket: state.socket(for: winner), remote: winner,
            sendToken: hostToken, recvToken: state.localToken)
        mux.attachDirect(rudp)
        let switched = await mux.verifyAndSwitch()
        DirectDebug.log("client.go verifyAndSwitch → \(switched)")
        finished = true
        if !switched { state.sockets.close() }
        self.state = nil
        return switched
    }

    public func abandon() async {
        guard !finished else { return }
        finished = true
        await engine?.stop()
        engine = nil
        state?.sockets.close()
        state = nil
    }
}

/// The host's direct-path replies, decoupled from any one event-stream shape so the Mac peer
/// client and the iPhone's `HostConnection` can both feed it: the consumer awaits
/// `next(timeout:)`, the connection's event loop calls `deliver(_:)`.
public final class DirectReplyBox: @unchecked Sendable {
    public enum HostReply: Sendable {
        case answer(DirectPunchOffer)
        case go
        case decline(String)
    }

    private let lock = NSLock()
    private var buffer: [HostReply] = []
    private var waiter: CheckedContinuation<HostReply?, Never>?
    private var finished = false

    public init() {}

    public func deliver(_ reply: HostReply) {
        lock.lock()
        if let waiter {
            self.waiter = nil
            lock.unlock()
            waiter.resume(returning: reply)
            return
        }
        buffer.append(reply)
        lock.unlock()
    }

    public func finish() {
        lock.lock()
        finished = true
        let waiter = self.waiter
        self.waiter = nil
        lock.unlock()
        waiter?.resume(returning: nil)
    }

    public func next(timeout: Duration) async -> HostReply? {
        await withCheckedContinuation { (cont: CheckedContinuation<HostReply?, Never>) in
            lock.lock()
            if !buffer.isEmpty {
                let first = buffer.removeFirst()
                lock.unlock()
                cont.resume(returning: first)
                return
            }
            if finished {
                lock.unlock()
                cont.resume(returning: nil)
                return
            }
            waiter = cont
            lock.unlock()
            Task { [weak self] in
                try? await ContinuousClock().sleep(for: timeout)
                self?.timeoutWaiter()
            }
        }
    }

    private func timeoutWaiter() {
        lock.lock()
        let waiter = self.waiter
        self.waiter = nil
        lock.unlock()
        waiter?.resume(returning: nil)
    }
}

/// The client-side upgrade session driver: offer → punch on answer → select → dial on go →
/// switch, with a small retry budget. Runs beside the session's event loop; a live session
/// is never blocked on it, and every failure just leaves the session on relay as it was.
public enum DirectUpgradeDriver {
    public static func run(
        mux: MultipathFrameChannel,
        attempts: Int = 3,
        retryDelay: Duration = .seconds(60),
        replyTimeout: Duration = .seconds(20),
        replies: DirectReplyBox,
        send: @escaping @Sendable (ClientMsg) async -> Void
    ) async -> Bool {
        var attempt = 0
        while attempt < attempts, !Task.isCancelled, mux.canAttachDirect {
            attempt += 1
            let upgrader = DirectClientUpgrader(mux: mux)
            guard let offer = await upgrader.makeOffer() else { return false }
            await send(.directOffer(offer))

            guard case .answer(let hostOffer) = await replies.next(timeout: replyTimeout) else {
                await upgrader.abandon()
                guard !Task.isCancelled else { return false }
                try? await Task.sleep(for: retryDelay)
                continue
            }
            guard let selected = await upgrader.answer(hostOffer) else {
                guard !Task.isCancelled else { return false }
                try? await Task.sleep(for: retryDelay)
                continue
            }
            await send(.directSelect(selected))
            guard case .go = await replies.next(timeout: replyTimeout) else {
                await upgrader.abandon()
                guard !Task.isCancelled else { return false }
                try? await Task.sleep(for: retryDelay)
                continue
            }
            if await upgrader.go() { return true }
            try? await Task.sleep(for: retryDelay)
        }
        return false
    }
}
