import Foundation

/// The punch probe wire format (SYNC_PROTOCOL §3.3): tiny datagrams both sides fire at each
/// other's candidates simultaneously until a round-trip proves a path. Every datagram carries
/// the **receiver's** expected token — 16 secret bytes exchanged inside Noise — so validation
/// is proof the sender is the authenticated peer: an off-path host can't validate a path, and
/// a NAT's stray reflections fail the token check.
///
/// `[5-byte magic "NUCP1"][1-byte type][16-byte token]` — 22 bytes, unfragmentable anywhere.
public enum PunchWire {
    public static let magic = Data([0x4E, 0x55, 0x43, 0x50, 0x31]) // "NUCP1"
    public static let tokenLength = 16
    public static let datagramLength = 5 + 1 + tokenLength

    public enum Kind: UInt8 {
        /// "Can you hear me?" — sent on a cadence to every remote candidate.
        case punch = 0x01
        /// "I hear you" — replied to the *observed source* of a valid punch, which is the
        /// address that actually traversed the NAT (not necessarily any advertised candidate).
        case ack = 0x02
    }

    public static func makeToken() -> Data {
        Data((0..<tokenLength).map { _ in UInt8.random(in: .min ... .max) })
    }

    public static func encode(_ kind: Kind, token: Data) -> Data {
        precondition(token.count == tokenLength, "punch token must be \(tokenLength) bytes")
        var out = Data(capacity: datagramLength)
        out.append(magic)
        out.append(kind.rawValue)
        out.append(token)
        return out
    }

    public static func decode(_ datagram: Data) -> (kind: Kind, token: Data)? {
        guard datagram.count == datagramLength, datagram.prefix(5) == magic,
              let kind = Kind(rawValue: datagram[datagram.startIndex + 5])
        else { return nil }
        return (kind, Data(datagram.suffix(tokenLength)))
    }
}

/// Punch policy: attempt unless *both* ends are behind symmetric NATs (both mapped ports
/// are unpredictable ⇒ probes can't meet). One symmetric side still punches: the cone
/// side's advertised port is stable, and the symmetric side's probes to it open its own
/// return path. `unknown` is treated as punchable — probes are 22 bytes and the session
/// already lives on the relay, so optimism costs nothing.
public func punchWorthAttempting(local: NATClass, remote: NATClass) -> Bool {
    !(local == .symmetric && remote == .symmetric)
}

/// The deterministic core of one punch attempt — pure state machine, no sockets, no clock:
/// the driver (or a test, or the NAT sim harness) feeds `Input`s and applies `Output`s.
/// Both sides run the same machine; there are no roles at this layer.
public struct PunchSession: Sendable {
    public enum Input: Sendable {
        /// Cadence tick: (re)fire probes at every remote candidate.
        case tick
        /// A datagram arrived on one of our punch sockets.
        case datagram(Data, from: UDPEndpoint)
    }

    public enum Output: Sendable, Equatable {
        case send(Data, to: UDPEndpoint)
        /// A round trip completed: our punch reached the peer and its ack came back from
        /// `endpoint` — this is a dialable address for the QUIC leg.
        case validated(UDPEndpoint)
        /// The peer's punch reached us from `endpoint` (half-open). Not sufficient to dial —
        /// but the host side uses it to know punching is converging.
        case inboundValidated(UDPEndpoint)
    }

    /// The token we advertised inside Noise — inbound datagrams must carry it.
    private let localToken: Data
    /// The token the peer advertised — we stamp it on everything we send.
    private let remoteToken: Data
    private let remoteCandidates: [UDPEndpoint]
    private var validatedEndpoints = Set<UDPEndpoint>()
    private var inboundEndpoints = Set<UDPEndpoint>()

    public init(localToken: Data, remoteToken: Data, remoteCandidates: [DirectCandidate]) {
        self.localToken = localToken
        self.remoteToken = remoteToken
        // Preserve the offer's probing order (v6 → reflexive → local) but dedup: several
        // candidates can share an endpoint when interfaces overlap.
        var seen = Set<UDPEndpoint>()
        self.remoteCandidates = remoteCandidates.map(\.endpoint).filter { seen.insert($0).inserted }
    }

    /// Every dialable endpoint proven by a full round trip so far, in discovery order.
    public private(set) var validated: [UDPEndpoint] = []

    public mutating func handle(_ input: Input) -> [Output] {
        switch input {
        case .tick:
            // Keep probing even after validation: the peer may still be waiting on its own
            // round trip, and our probes keep the NAT bindings warm until the QUIC handoff.
            return remoteCandidates.map {
                .send(PunchWire.encode(.punch, token: remoteToken), to: $0)
            }

        case .datagram(let data, let from):
            guard let decoded = PunchWire.decode(data), decoded.token == localToken else {
                return [] // not a punch datagram, or not our authenticated peer — ignore
            }
            switch decoded.kind {
            case .punch:
                // Ack to the observed source: that address traversed the NAT, the
                // advertised candidate it corresponds to may not exist from our vantage.
                var outputs: [Output] = [
                    .send(PunchWire.encode(.ack, token: remoteToken), to: from),
                ]
                if inboundEndpoints.insert(from).inserted {
                    outputs.append(.inboundValidated(from))
                }
                return outputs
            case .ack:
                if validatedEndpoints.insert(from).inserted {
                    validated.append(from)
                    return [.validated(from)]
                }
                return []
            }
        }
    }
}

/// Drives a `PunchSession` over real sockets: ticks on a cadence, pumps inbound datagrams,
/// reports validations. Runs until cancelled or the budget lapses — deliberately *not*
/// stopping at first validation, because the peer needs our continued probes/acks to finish
/// its own round trip, and the coordinator (not this engine) decides when to hand off to QUIC.
public actor PunchEngine {
    public struct Update: Sendable {
        public let endpoint: UDPEndpoint
        /// True for a full round trip (dialable), false for inbound-only.
        public let roundTrip: Bool
    }

    public static let defaultCadence = Duration.milliseconds(200)
    public static let defaultBudget = Duration.seconds(10)

    private let sockets: DirectCandidateGatherer.Sockets
    private var session: PunchSession

    public init(
        sockets: DirectCandidateGatherer.Sockets,
        localToken: Data,
        remoteToken: Data,
        remoteCandidates: [DirectCandidate]
    ) {
        self.sockets = sockets
        self.session = PunchSession(
            localToken: localToken, remoteToken: remoteToken, remoteCandidates: remoteCandidates)
    }

    /// Probe until the budget lapses or the task is cancelled, streaming validations as
    /// they land. The stream finishes at the end of the attempt; `validated` afterwards
    /// holds every proven endpoint.
    public func run(
        cadence: Duration = PunchEngine.defaultCadence,
        budget: Duration = PunchEngine.defaultBudget
    ) -> AsyncStream<Update> {
        AsyncStream { continuation in
            let pump = Task {
                let clock = ContinuousClock()
                let deadline = clock.now + budget

                let ticker = Task {
                    while !Task.isCancelled, clock.now < deadline {
                        self.apply(self.handle(.tick), continuation: continuation)
                        try? await clock.sleep(for: cadence)
                    }
                }
                // One receive loop per socket family.
                await withTaskGroup(of: Void.self) { group in
                    for socket in [sockets.v4, sockets.v6].compactMap({ $0 }) {
                        group.addTask {
                            while !Task.isCancelled {
                                guard
                                    let packet = try? await socket.receive(
                                        deadline: deadline, clock: clock),
                                    !Task.isCancelled
                                else { break }
                                await self.pump(packet, continuation: continuation)
                            }
                        }
                    }
                }
                ticker.cancel()
                continuation.finish()
            }
            // `AsyncStream` iteration doesn't propagate task cancellation to `pump`, so a
            // consumer that stops early (a `break`, or cancelling its consuming task) would
            // leave `pump`'s receive loops reading the shared socket for the whole budget —
            // and RUDP, which adopts that same socket after the punch, would never see a
            // datagram. `stop()` is the explicit off switch the coordinators call; the
            // onTermination is the backstop for the stream finishing on its own.
            self.pumpTask = pump
            continuation.onTermination = { _ in pump.cancel() }
        }
    }

    private var pumpTask: Task<Void, Never>?

    /// Halt all probing and release the socket immediately (the coordinators call this once
    /// a pair is chosen, before handing the socket to RUDP).
    public func stop() {
        pumpTask?.cancel()
        pumpTask = nil
    }

    /// Every round-trip-proven endpoint so far, in validation order.
    public var validated: [UDPEndpoint] { session.validated }

    private func pump(
        _ packet: (data: Data, from: UDPEndpoint), continuation: AsyncStream<Update>.Continuation
    ) {
        apply(handle(.datagram(packet.data, from: packet.from)), continuation: continuation)
    }

    private func handle(_ input: PunchSession.Input) -> [PunchSession.Output] {
        session.handle(input)
    }

    private func apply(
        _ outputs: [PunchSession.Output], continuation: AsyncStream<Update>.Continuation
    ) {
        for output in outputs {
            switch output {
            case .send(let data, let to):
                let socket = to.isIPv6 ? sockets.v6 : sockets.v4
                try? socket?.send(data, to: to)
            case .validated(let endpoint):
                continuation.yield(Update(endpoint: endpoint, roundTrip: true))
            case .inboundValidated(let endpoint):
                continuation.yield(Update(endpoint: endpoint, roundTrip: false))
            }
        }
    }
}
