import Foundation

/// A reliable, ordered byte pipe over the *already-punched* UDP socket (Covalence direct,
/// SYNC_PROTOCOL §3.3) — the direct leg's transport once a hole is open. It presents the
/// same `FrameChannel` surface as every other transport: `WireFraming` length-prefixed
/// frames on top of a reliable stream, with Noise (the security boundary) and the mux above
/// it, all transport-blind.
///
/// **Why not QUIC.** The first cut dialed `NWProtocolQUIC` over the punched port. Network
/// .framework's QUIC couples the data path to Apple-only APIs (so Linux `nucleicd` couldn't
/// use it) and to a fragile close-punch-socket-then-rebind-the-port handoff. RUDP keeps the
/// *live* punched socket, is pure POSIX (Linux included), and needs no TLS identity —
/// Noise already end-to-end-encrypts every frame, and the punch token authenticates the
/// peer, so a plaintext reliable layer underneath is exactly right.
///
/// **Reliability.** Go-back-N: a sliding send window of sequenced datagrams, cumulative
/// ACKs, and timeout retransmission from the cumulative-ack point; the receiver delivers
/// strictly in order and drops out-of-order datagrams (the sender resends them). Simple and
/// provably in-order — which is what the Noise nonce discipline above requires. Every
/// datagram carries the peer's 16-byte punch token, so an off-path host that doesn't know
/// the token (exchanged only inside Noise) can't inject into the stream.
///
/// Datagram: `[4 magic "NRUD"][16 token][1 type][8 seq BE][2 len BE][payload]`.
public final class RUDPFrameChannel: FrameChannel, @unchecked Sendable {
    static let magic = Data("NRUD".utf8)
    static let headerLength = 4 + 16 + 1 + 8 + 2
    /// Conservative payload cap so a datagram clears the smallest plausible path MTU
    /// (1280 IPv6 minimum − headers) without IP fragmentation.
    static let maxPayload = 1100
    static let windowSize = 256
    static let retransmitInterval = Duration.milliseconds(200)
    static let keepaliveInterval = Duration.seconds(15)
    /// No progress for this long ⇒ the path is dead; end the stream so the mux fails back.
    static let deadPathTimeout = Duration.seconds(30)

    enum Kind: UInt8 {
        case data = 0x01
        case ack = 0x02
        /// Empty keepalive so an idle path refreshes NAT bindings and proves liveness.
        case ping = 0x03
    }

    public let peerDescription: String
    public var transport: SyncTransportHint? { .direct }

    private let socket: UDPSocket
    /// Where we send — seeded from the punch-validated endpoint, then **latched** to the
    /// source of the first token-valid inbound datagram. The punch may validate an asymmetric
    /// address pair (a multi-homed host, or a NAT that rebinds), so the address that actually
    /// delivers is authoritative; the token (exchanged only inside Noise) is what authorizes
    /// adopting it. `lock`-guarded.
    private var remote: UDPEndpoint
    private var remoteLatched = false
    /// The token we stamp on our datagrams (the peer's expected token).
    private let sendToken: Data
    /// The token inbound datagrams must carry (our advertised token).
    private let recvToken: Data
    private let clock = ContinuousClock()

    private let lock = NSLock()
    private var closed = false

    // Send state (our bytes → peer).
    private var outBuffer = Data()               // unframed bytes awaiting chunking
    private var nextSeq: UInt64 = 0
    private var sendBase: UInt64 = 0             // oldest unacked seq
    private var window: [(seq: UInt64, datagram: Data)] = []
    private var lastAckAdvance: ContinuousClock.Instant

    // Receive state (peer's bytes → us).
    private var rcvNext: UInt64 = 0
    private let accumulator = FrameAccumulator()

    private var outStream: AsyncStream<Data>
    private var outContinuation: AsyncStream<Data>.Continuation?
    private var pumpTask: Task<Void, Never>?
    private var timerTask: Task<Void, Never>?

    public init(socket: UDPSocket, remote: UDPEndpoint, sendToken: Data, recvToken: Data) {
        self.socket = socket
        self.remote = remote
        self.sendToken = sendToken
        self.recvToken = recvToken
        self.peerDescription = "direct:\(remote.description)"
        self.lastAckAdvance = clock.now
        var cont: AsyncStream<Data>.Continuation?
        self.outStream = AsyncStream(bufferingPolicy: .unbounded) { cont = $0 }
        self.outContinuation = cont
        start()
    }

    private func start() {
        pumpTask = Task { [weak self] in await self?.receiveLoop() }
        timerTask = Task { [weak self] in await self?.timerLoop() }
    }

    public func frames() -> AsyncStream<Data> { outStream }

    public func send(_ frame: Data) {
        lock.lock()
        guard !closed else { lock.unlock(); return }
        outBuffer.append(WireFraming.frame(frame))
        let datagrams = fillWindowLocked()
        let dest = remote
        lock.unlock()
        for d in datagrams { try? socket.send(d, to: dest) }
    }

    public func close() {
        lock.lock()
        guard !closed else { lock.unlock(); return }
        closed = true
        let cont = outContinuation
        outContinuation = nil
        lock.unlock()
        pumpTask?.cancel()
        timerTask?.cancel()
        cont?.finish()
    }

    // MARK: Sender

    /// Chunk `outBuffer` into datagrams while the window has room; returns the new datagrams
    /// to send. Caller holds `lock`.
    private func fillWindowLocked() -> [Data] {
        var out: [Data] = []
        while window.count < Self.windowSize, !outBuffer.isEmpty {
            let take = min(Self.maxPayload, outBuffer.count)
            let chunk = outBuffer.prefix(take)
            outBuffer.removeFirst(take)
            let seq = nextSeq
            nextSeq += 1
            let datagram = Self.encode(.data, token: sendToken, seq: seq, payload: Data(chunk))
            window.append((seq, datagram))
            out.append(datagram)
        }
        return out
    }

    private func timerLoop() async {
        while !Task.isCancelled {
            try? await clock.sleep(for: Self.retransmitInterval)
            guard !Task.isCancelled else { return }
            let tick = timerTick()
            if tick.dead { close(); return }
            for d in tick.resend { try? socket.send(d, to: tick.dest) }
            if let ping = tick.ping { try? socket.send(ping, to: tick.dest) }
        }
    }

    /// Sync (lockable) half of one timer tick — `NSLock` is unavailable from async contexts,
    /// so the async loop hops through here.
    private func timerTick() -> (dead: Bool, resend: [Data], ping: Data?, dest: UDPEndpoint) {
        lock.lock()
        defer { lock.unlock() }
        guard !closed else { return (false, [], nil, remote) }
        let idle = clock.now - lastAckAdvance
        if !window.isEmpty {
            if idle > Self.deadPathTimeout { return (true, [], nil, remote) }
            // Go-back-N: resend the whole outstanding window from sendBase.
            return (false, window.map(\.datagram), nil, remote)
        }
        if idle > Self.keepaliveInterval {
            return (false, [], Self.encode(.ping, token: sendToken, seq: nextSeq, payload: Data()), remote)
        }
        return (false, [], nil, remote)
    }

    // MARK: Receiver

    private func receiveLoop() async {
        while !Task.isCancelled {
            guard let packet = try? await socket.receive(deadline: clock.now + Self.keepaliveInterval)
            else {
                if Task.isCancelled { return }
                continue // deadline lapsed with no datagram — loop (the timer handles liveness)
            }
            handle(packet.data, from: packet.from)
        }
    }

    private func handle(_ datagram: Data, from source: UDPEndpoint) {
        // The token (Noise-exchanged) authorizes the datagram; the source address does not.
        guard let (kind, token, seq, payload) = Self.decode(datagram), token == recvToken else {
            return
        }
        // Latch our send target onto the address that actually delivered an authenticated
        // datagram — the punch may validate an asymmetric pair (multi-homed host / NAT
        // rebind), so this is the real working path.
        lock.lock()
        if !remoteLatched {
            remote = source
            remoteLatched = true
        }
        let dest = remote
        lock.unlock()

        switch kind {
        case .ping:
            // Liveness only; ack our current expectation so the peer's keepalive advances too.
            sendAck()
        case .ack:
            lock.lock()
            if closed { lock.unlock(); return }
            // `seq` is the peer's next-expected — everything below it is acked.
            let before = sendBase
            while let first = window.first, first.seq < seq {
                window.removeFirst()
            }
            sendBase = max(sendBase, seq)
            if sendBase > before { lastAckAdvance = clock.now }
            let more = fillWindowLocked()
            lock.unlock()
            for d in more { try? socket.send(d, to: dest) }
        case .data:
            var deliver: [Data] = []
            lock.lock()
            if closed { lock.unlock(); return }
            if seq == rcvNext {
                rcvNext += 1
                lastAckAdvance = clock.now
                if let frames = try? accumulator.push(payload) {
                    deliver = frames
                } else {
                    // Oversize/garbage frame — kill the leg; the mux fails back.
                    lock.unlock()
                    close()
                    return
                }
            }
            // In-order accepted, duplicate, or ahead-of-order: always (re-)ACK rcvNext so a
            // lost ACK or a go-back-N resend converges.
            let cont = outContinuation
            lock.unlock()
            for frame in deliver { cont?.yield(frame) }
            sendAck()
        }
    }

    private func sendAck() {
        lock.lock()
        let next = rcvNext
        let token = sendToken
        let dest = remote
        lock.unlock()
        try? socket.send(Self.encode(.ack, token: token, seq: next, payload: Data()), to: dest)
    }

    // MARK: Wire

    static func encode(_ kind: Kind, token: Data, seq: UInt64, payload: Data) -> Data {
        precondition(token.count == 16, "rudp token must be 16 bytes")
        var out = Data(capacity: headerLength + payload.count)
        out.append(magic)
        out.append(token)
        out.append(kind.rawValue)
        out.append(contentsOf: withUnsafeBytes(of: seq.bigEndian, Array.init))
        let len = UInt16(min(payload.count, Int(UInt16.max)))
        out.append(UInt8(len >> 8))
        out.append(UInt8(len & 0xFF))
        out.append(payload)
        return out
    }

    static func decode(_ datagram: Data) -> (kind: Kind, token: Data, seq: UInt64, payload: Data)? {
        guard datagram.count >= headerLength, datagram.prefix(4) == magic else { return nil }
        let base = datagram.startIndex
        let token = datagram.subdata(in: base + 4 ..< base + 20)
        guard let kind = Kind(rawValue: datagram[base + 20]) else { return nil }
        var seq: UInt64 = 0
        for i in 0..<8 { seq = (seq << 8) | UInt64(datagram[base + 21 + i]) }
        let len = Int(datagram[base + 29]) << 8 | Int(datagram[base + 30])
        let payloadStart = base + headerLength
        guard datagram.count >= headerLength + len else { return nil }
        let payload = datagram.subdata(in: payloadStart ..< payloadStart + len)
        return (kind, token, seq, payload)
    }
}
