import Foundation

#if canImport(CryptoKit)
import CryptoKit
#else
import Crypto
#endif

/// One Noise session, two byte paths (SYNC_PROTOCOL §3.3): wraps the relay `FrameChannel`
/// a session was established on and, when a hole-punched QUIC leg comes up, migrates the
/// session's frames onto it — keeping the relay leg open, idle, as the instant failback.
///
/// **Why this is delicate**: Noise nonces are implicit counters, so the receiver must
/// decrypt frames in exactly the order they were sealed, and no sealed frame may ever be
/// lost while the session lives (`FrameChannel.send` doc). The mux therefore:
///   - keeps ONE active outbound path per direction, flipped by a `SWITCH` marker sent as
///     the last frame on the old path (both legs are individually ordered, so releasing
///     "old path until the marker, then the new path" preserves total order);
///   - sequences every frame it sends on the direct leg (`[8-byte seq][frame]`) and holds
///     it in a retransmit buffer until the peer's cumulative `ACK`;
///   - on direct-path death, runs a `RESYNC` exchange over the relay leg: each side closes
///     QUIC, *drains what was already delivered*, reports the last direct seq it processed,
///     and the peer retransmits everything after that — raw, in order, on the relay. The
///     gap is filled exactly: no duplicate, no loss, no re-handshake.
///
/// **Locking discipline**: every underlying `send` and every `yield` to the outbound stream
/// happens *inside* the mux lock. That single rule is what makes the ordering guarantees
/// hold — a SWITCH marker can never overtake a data frame, and releases from the two pump
/// tasks can never interleave. Underlying sends are non-blocking enqueues and
/// `AsyncStream.Continuation.yield` never suspends, so holding the lock across them is safe.
///
/// **Passthrough until armed**: constructed around any channel it is a zero-cost forwarder;
/// only `arm(bindingKey:)` — called after hello/welcome, with a key derived from the Noise
/// handshake hash — turns on control-frame interpretation. Sessions that never upgrade
/// never change behavior. Control frames carry an HMAC under that binding key: the relay
/// (or any on-path party) can drop them (it can already drop anything) but cannot forge a
/// switch or a resync.
///
/// **If the relay leg dies**, upgraded or not, the mux finishes its stream and the session
/// tears down exactly like any dead connection today — the relay is the session's spine.
public final class MultipathFrameChannel: FrameChannel, @unchecked Sendable {
    // MARK: Control wire format

    /// `[16-byte magic][1 type][8 generation BE][8 value BE][32 HMAC-SHA256]` = 65 bytes.
    /// On the relay leg a control frame is any frame starting with the magic (a Noise
    /// ciphertext colliding is 2⁻¹²⁸ — accepted); on the direct leg control frames ride
    /// under the reserved seq sentinel instead.
    enum Control {
        static let magic = Data("NUCLEIC-MUX-0001".utf8) // exactly 16 bytes
        static let length = 16 + 1 + 8 + 8 + 32

        enum Kind: UInt8 {
            /// Relay leg, per direction: "my frames now continue on the direct leg,
            /// starting at seq `value`". The flip point for the receiver's release order.
            case switchToDirect = 0x01
            /// Direct leg: cumulative "I have processed direct frames through seq `value`" —
            /// lets the sender trim its retransmit buffer.
            case ack = 0x02
            /// Direct leg: liveness + peer-authentication probe (`value` = random nonce).
            case probe = 0x03
            case probeAck = 0x04
            /// Relay leg: "the direct leg is dead; I processed your direct frames through
            /// seq `value` — retransmit the rest on the relay and answer with your count."
            case resync = 0x05
            case resyncAck = 0x06
        }

        static func encode(_ kind: Kind, generation: UInt64, value: UInt64, key: SymmetricKey) -> Data {
            var frame = Data(capacity: length)
            frame.append(magic)
            frame.append(kind.rawValue)
            frame.append(contentsOf: withUnsafeBytes(of: generation.bigEndian, Array.init))
            frame.append(contentsOf: withUnsafeBytes(of: value.bigEndian, Array.init))
            let mac = HMAC<SHA256>.authenticationCode(for: frame, using: key)
            frame.append(contentsOf: mac)
            return frame
        }

        static func decode(_ frame: Data, key: SymmetricKey) -> (kind: Kind, generation: UInt64, value: UInt64)? {
            guard frame.count == length, frame.prefix(16) == magic else { return nil }
            let body = frame.prefix(length - 32)
            let mac = frame.suffix(32)
            guard HMAC<SHA256>.isValidAuthenticationCode(mac, authenticating: body, using: key),
                  let kind = Kind(rawValue: frame[frame.startIndex + 16])
            else { return nil }
            let generation = frame.readUInt64(at: 17)
            let value = frame.readUInt64(at: 25)
            return (kind, generation, value)
        }

        static func isControl(_ frame: Data) -> Bool {
            frame.count == length && frame.prefix(16) == magic
        }
    }

    /// Direct-leg framing: `[8-byte seq BE][payload]`; `seq == .max` marks a control frame.
    static let controlSeqSentinel = UInt64.max

    // MARK: State (all mutable state behind `lock`)

    private enum OutboundPath {
        case relay
        case direct
        /// Mid-failback: new sends queue up until the RESYNC exchange finishes.
        case failingBack
    }

    private let primary: any FrameChannel
    private let lock = NSLock()

    private var bindingKey: SymmetricKey?
    private var direct: (any FrameChannel)?
    private var directPump: Task<Void, Never>?
    private var outbound: OutboundPath = .relay
    private var inboundDirectActive = false
    private var generation: UInt64 = 0
    private var closed = false

    // Outbound direct bookkeeping (our frames → peer).
    private var nextDirectSeq: UInt64 = 1
    private var retransmit: [(seq: UInt64, frame: Data)] = []
    private var retransmitBytes = 0
    private var pendingDuringFailback: [Data] = []
    private var lastAckProgress = ContinuousClock().now

    // Inbound direct bookkeeping (peer's frames → us).
    private var lastDirectSeqReceived: UInt64 = 0
    private var directBuffer: [(seq: UInt64, frame: Data)] = []
    private var directInboundSinceAck = 0

    // Verify / resync rendezvous.
    private var probeNonce: UInt64?
    private var probeContinuation: CheckedContinuation<Bool, Never>?
    private var resyncContinuation: CheckedContinuation<UInt64?, Never>?
    /// A peer count that arrived before our own failback reached the waiting phase
    /// (simultaneous detection) — consumed instead of waiting.
    private var pendingPeerResync: UInt64?
    private var watchdog: Task<Void, Never>?

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

    private var _onPathChanged: (@Sendable (SyncTransportHint) -> Void)?

    /// UI/presence hook: fires with `.direct` when frames start flowing directly and with
    /// the primary's transport again after a failback. Set before the upgrade starts.
    public var onPathChanged: (@Sendable (SyncTransportHint) -> Void)? {
        get { lock.withLock { _onPathChanged } }
        set { lock.withLock { _onPathChanged = newValue } }
    }

    /// Retransmit-buffer ceiling: past this the direct path isn't keeping up (or its acks
    /// are blackholed) and the mux proactively fails back — the buffer is the loss-proofing
    /// and must not grow unboundedly.
    private let retransmitCap: Int

    public init(primary: any FrameChannel, retransmitCap: Int = 8 * 1024 * 1024) {
        self.primary = primary
        self.retransmitCap = retransmitCap
        var continuation: AsyncStream<Data>.Continuation?
        self.outStream = AsyncStream { continuation = $0 }
        self.outContinuation = continuation
        startPrimaryPump()
    }

    // MARK: FrameChannel (the surface ConnectionHandler / SyncClient see)

    public var peerDescription: String { primary.peerDescription }

    /// Deliberately the *primary's* transport even while direct is active: host logic keys
    /// decisions off the accepted transport (`.relay`), and the session's spine is still the
    /// relay. The live data path surfaces via `activePath` / `onPathChanged` instead.
    public var transport: SyncTransportHint? { primary.transport }

    /// What's actually carrying frames right now — presence/UI truth.
    public var activePath: SyncTransportHint {
        lock.withLock { outbound == .direct ? .direct : (primary.transport ?? .relay) }
    }

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

    public func send(_ frame: Data) {
        var overCap = false
        lock.withLock {
            guard !closed else { return }
            switch outbound {
            case .relay:
                primary.send(frame)
            case .failingBack:
                pendingDuringFailback.append(frame)
            case .direct:
                let seq = nextDirectSeq
                nextDirectSeq += 1
                retransmit.append((seq, frame))
                retransmitBytes += frame.count
                overCap = retransmitBytes > retransmitCap
                direct?.send(Self.sequenced(seq: seq, frame: frame))
            }
        }
        if overCap { scheduleFailback() }
    }

    public func close() {
        let teardown: ((any FrameChannel)?, Task<Void, Never>?, Task<Void, Never>?)? = lock.withLock {
            guard !closed else { return nil }
            closed = true
            return (direct, directPump, watchdog)
        }
        guard let (directChannel, pump, wd) = teardown else { return }
        wd?.cancel()
        pump?.cancel()
        directChannel?.close()
        primary.close()
        resumeProbe(false)
        resumeResync(nil)
        outContinuation?.finish()
    }

    // MARK: Upgrade lifecycle (driven by the Direct coordinators)

    /// Turn on control-frame interpretation. `bindingKey` must be derived from the session's
    /// Noise handshake hash on both ends (`Self.bindingKey(handshakeHash:)`) — it is what
    /// makes SWITCH/RESYNC unforgeable by the relay.
    public func arm(bindingKey: SymmetricKey) {
        lock.withLock { self.bindingKey = bindingKey }
    }

    public static func bindingKey(handshakeHash: Data) -> SymmetricKey {
        HKDF<SHA256>.deriveKey(
            inputKeyMaterial: SymmetricKey(data: handshakeHash),
            info: Data("nucleic-direct-mux-v1".utf8), outputByteCount: 32)
    }

    /// Whether an upgrade can start: armed, open, no direct leg attached, riding the relay.
    public var canAttachDirect: Bool {
        lock.withLock { bindingKey != nil && !closed && direct == nil && outbound == .relay }
    }

    /// Adopt a freshly-punched direct leg. Frames it delivers are buffered (not released)
    /// until the peer's `SWITCH` marker arrives on the relay leg; our own outbound flips
    /// only after `verifyAndSwitch()`.
    public func attachDirect(_ channel: any FrameChannel) {
        let accepted: Bool = lock.withLock {
            guard bindingKey != nil, !closed, direct == nil, outbound == .relay else { return false }
            direct = channel
            generation += 1
            return true
        }
        guard accepted else {
            channel.close()
            return
        }
        watchdog?.cancel()
        startDirectPump(channel)
        startWatchdog()
    }

    /// Prove the attached direct leg carries our authenticated peer (PROBE/PROBE_ACK under
    /// the binding key), then flip our outbound onto it — emitting the SWITCH marker as the
    /// last frame we put on the relay leg. Returns false (and detaches the leg) on timeout.
    public func verifyAndSwitch(timeout: Duration = .seconds(5)) async -> Bool {
        let nonce = UInt64.random(in: .min ... .max)
        let sent: Bool = lock.withLock {
            guard let key = bindingKey, let channel = direct, outbound == .relay, !closed,
                  probeContinuation == nil
            else { return false }
            probeNonce = nonce
            channel.send(
                Self.sequenced(
                    seq: Self.controlSeqSentinel,
                    frame: Control.encode(.probe, generation: generation, value: nonce, key: key)))
            return true
        }
        guard sent else { return false }

        let verified = await withCheckedContinuation { (cont: CheckedContinuation<Bool, Never>) in
            let alreadyDone: Bool = lock.withLock {
                // The ack could have raced in between send and here.
                guard probeNonce == nonce else { return true }
                probeContinuation = cont
                return false
            }
            if alreadyDone {
                cont.resume(returning: true)
                return
            }
            Task { [weak self] in
                try? await ContinuousClock().sleep(for: timeout)
                self?.resumeProbe(false)
            }
        }
        guard verified else {
            detachDirect()
            return false
        }

        // Flip: SWITCH is the last frame we put on the relay leg; everything after rides
        // direct. Marker + flag flip under one lock hold, so no send can slip between them.
        var notify: (@Sendable (SyncTransportHint) -> Void)?
        let flipped: Bool = lock.withLock {
            guard let key = bindingKey, outbound == .relay, direct != nil, !closed else {
                return false
            }
            primary.send(
                Control.encode(
                    .switchToDirect, generation: generation, value: nextDirectSeq, key: key))
            outbound = .direct
            lastAckProgress = ContinuousClock().now
            notify = _onPathChanged
            return true
        }
        if flipped { notify?(.direct) }
        return flipped
    }

    /// Coordinator-triggered failback (e.g. a network-path change): same lossless RESYNC
    /// exchange as an organic direct-leg death.
    public func failbackToRelay() {
        scheduleFailback()
    }

    // MARK: Pumps

    private func startPrimaryPump() {
        primaryPump = Task { [weak self] in
            guard let stream = self?.primary.frames() else { return }
            for await frame in stream {
                guard let self else { return }
                self.handlePrimaryFrame(frame)
            }
            // Relay leg finished ⇒ the session's spine is gone. Tear down like any dead
            // connection — even a live direct leg has lost its control plane.
            self?.close()
        }
    }

    private func handlePrimaryFrame(_ frame: Data) {
        let key = lock.withLock { bindingKey }
        if let key, Control.isControl(frame) {
            guard let control = Control.decode(frame, key: key) else { return } // bad HMAC: drop
            handleControl(control)
            return
        }
        lock.withLock {
            guard !closed else { return }
            outContinuation?.yield(frame)
        }
    }

    private func startDirectPump(_ channel: any FrameChannel) {
        directPump = Task { [weak self] in
            for await raw in channel.frames() {
                guard let self, let (seq, frame) = Self.desequence(raw) else { continue }
                self.handleDirectFrame(seq: seq, frame: frame)
            }
            // Drained: everything the leg delivered has been processed —
            // `lastDirectSeqReceived` is final, which is exactly what RESYNC reports.
            self?.directLegFinished()
        }
    }

    private func handleDirectFrame(seq: UInt64, frame: Data) {
        if seq == Self.controlSeqSentinel {
            let key = lock.withLock { bindingKey }
            guard let key, let control = Control.decode(frame, key: key) else { return }
            handleControl(control)
            return
        }

        var gap = false
        lock.withLock {
            guard !closed else { return }
            if inboundDirectActive {
                guard seq == lastDirectSeqReceived + 1 else {
                    // A gap on an ordered stream means the leg is corrupt — fail back.
                    gap = true
                    return
                }
                lastDirectSeqReceived = seq
                outContinuation?.yield(frame)
                directInboundSinceAck += 1
                if directInboundSinceAck >= 32, let key = bindingKey, let channel = direct {
                    directInboundSinceAck = 0
                    channel.send(
                        Self.sequenced(
                            seq: Self.controlSeqSentinel,
                            frame: Control.encode(
                                .ack, generation: generation, value: lastDirectSeqReceived, key: key)))
                }
            } else {
                // Peer already switched; our release point (its SWITCH marker) hasn't
                // arrived on the relay leg yet. Hold — releasing early would jump the relay
                // frames still ahead of the marker and desync the Noise nonce order.
                directBuffer.append((seq, frame))
            }
        }
        if gap { scheduleFailback() }
    }

    // MARK: Control handling

    private func handleControl(_ control: (kind: Control.Kind, generation: UInt64, value: UInt64)) {
        switch control.kind {
        case .switchToDirect:
            // Peer's frames continue on the direct leg. Release anything the direct pump
            // already buffered, in order, then release live.
            lock.withLock {
                guard !closed else { return }
                inboundDirectActive = true
                directBuffer.sort { $0.seq < $1.seq }
                for (seq, frame) in directBuffer where seq == lastDirectSeqReceived + 1 {
                    lastDirectSeqReceived = seq
                    outContinuation?.yield(frame)
                }
                directBuffer.removeAll()
            }

        case .ack:
            lock.withLock {
                while let first = retransmit.first, first.seq <= control.value {
                    retransmitBytes -= first.frame.count
                    retransmit.removeFirst()
                }
                lastAckProgress = ContinuousClock().now
            }

        case .probe:
            lock.withLock {
                guard let key = bindingKey, let channel = direct else { return }
                channel.send(
                    Self.sequenced(
                        seq: Self.controlSeqSentinel,
                        frame: Control.encode(
                            .probeAck, generation: control.generation, value: control.value, key: key)))
            }

        case .probeAck:
            let matches = lock.withLock { probeNonce == control.value }
            if matches { resumeProbe(true) }

        case .resync:
            // Peer detected direct death. If we're mid-failback ourselves (simultaneous
            // detection) its RESYNC carries the same count its RESYNC_ACK would — feed the
            // waiter (or stash it if our failback hasn't reached the waiting phase yet).
            let respondAsFollower: Bool = lock.withLock {
                if outbound == .failingBack || closed {
                    if let cont = resyncContinuation {
                        resyncContinuation = nil
                        cont.resume(returning: control.value)
                    } else {
                        pendingPeerResync = control.value
                    }
                    return false
                }
                return true
            }
            if respondAsFollower {
                Task { [weak self] in await self?.performFailback(respondingTo: control.value) }
            }

        case .resyncAck:
            resumeResync(control.value)
        }
    }

    // MARK: Failback (the lossless RESYNC exchange)

    private func scheduleFailback() {
        Task { [weak self] in await self?.performFailback(respondingTo: nil) }
    }

    private func directLegFinished() {
        // Organic death (idle timeout, peer close, network drop). If we're already failing
        // back this is the expected drain completion; otherwise we detected it first.
        let shouldInitiate = lock.withLock {
            direct != nil && outbound != .failingBack && !closed
        }
        if shouldInitiate { scheduleFailback() }
    }

    /// Both roles of the RESYNC exchange. `peerReceived == nil` ⇒ we detected death and
    /// initiate (send RESYNC, await the answer); non-nil ⇒ the peer initiated and this is
    /// its received-count (we answer with RESYNC_ACK). Either way: close + drain our direct
    /// leg first, so both counts are final before anyone retransmits.
    private func performFailback(respondingTo peerReceived: UInt64?) async {
        // Phase 1: freeze outbound, close the direct leg.
        let setup: (key: SymmetricKey, gen: UInt64, channel: any FrameChannel, pump: Task<Void, Never>?)? =
            lock.withLock {
                guard let key = bindingKey, let channel = direct, !closed, outbound != .failingBack
                else { return nil }
                outbound = .failingBack
                return (key, generation, channel, directPump)
            }
        guard let setup else {
            // Already failing back (or no direct leg) — if the peer's count just arrived,
            // hand it to the waiting exchange.
            if let peerReceived { deliverPeerResync(peerReceived) }
            return
        }
        watchdog?.cancel()
        setup.channel.close()
        // Drain: the pump finishes once the closed channel's stream ends; after this,
        // `lastDirectSeqReceived` is final.
        await setup.pump?.value

        // Phase 2: exchange counts over the relay leg.
        var peerCount: UInt64? = peerReceived
        if peerReceived != nil {
            lock.withLock {
                primary.send(
                    Control.encode(
                        .resyncAck, generation: setup.gen, value: lastDirectSeqReceived, key: setup.key))
            }
        } else {
            peerCount = await withCheckedContinuation { (cont: CheckedContinuation<UInt64?, Never>) in
                let immediate: UInt64?? = lock.withLock {
                    if let stashed = pendingPeerResync {
                        pendingPeerResync = nil
                        return .some(stashed)
                    }
                    resyncContinuation = cont
                    primary.send(
                        Control.encode(
                            .resync, generation: setup.gen, value: lastDirectSeqReceived, key: setup.key))
                    return nil
                }
                if let immediate {
                    cont.resume(returning: immediate)
                    return
                }
                Task { [weak self] in
                    try? await ContinuousClock().sleep(for: .seconds(10))
                    self?.resumeResync(nil)
                }
            }
        }

        guard let peerCount else {
            // No answer on the relay leg within budget ⇒ the spine is effectively dead too.
            close()
            return
        }

        // Phase 3: fill the peer's gap — everything it hadn't processed, raw, in order —
        // then release the sends that queued during the exchange, then go live on relay.
        var notify: (@Sendable (SyncTransportHint) -> Void)?
        var restoredPath: SyncTransportHint = .relay
        lock.withLock {
            guard !closed else { return }
            for entry in retransmit where entry.seq > peerCount {
                primary.send(entry.frame)
            }
            retransmit.removeAll()
            retransmitBytes = 0
            for frame in pendingDuringFailback { primary.send(frame) }
            pendingDuringFailback.removeAll()
            direct = nil
            directPump = nil
            inboundDirectActive = false
            directBuffer.removeAll()
            directInboundSinceAck = 0
            outbound = .relay
            notify = _onPathChanged
            restoredPath = primary.transport ?? .relay
        }
        notify?(restoredPath)
    }

    private func deliverPeerResync(_ value: UInt64) {
        let cont: CheckedContinuation<UInt64?, Never>? = lock.withLock {
            if let cont = resyncContinuation {
                resyncContinuation = nil
                return cont
            }
            pendingPeerResync = value
            return nil
        }
        cont?.resume(returning: value)
    }

    private func detachDirect() {
        let teardown: ((any FrameChannel)?, Task<Void, Never>?) = lock.withLock {
            let channel = direct
            let pump = directPump
            direct = nil
            directPump = nil
            inboundDirectActive = false
            directBuffer.removeAll()
            return (channel, pump)
        }
        watchdog?.cancel()
        teardown.1?.cancel()
        teardown.0?.close()
    }

    /// ACK starvation guard: while our outbound rides the direct leg, unacked frames older
    /// than the threshold mean the leg (or its return path) is blackholed even though the
    /// transport still reports itself alive — fail back proactively.
    private func startWatchdog() {
        watchdog = Task { [weak self] in
            let clock = ContinuousClock()
            while !Task.isCancelled {
                try? await clock.sleep(for: .seconds(5))
                guard let self else { return }
                let starving = self.lock.withLock {
                    self.outbound == .direct && !self.retransmit.isEmpty
                        && clock.now - self.lastAckProgress > .seconds(15)
                }
                if starving {
                    self.scheduleFailback()
                    return
                }
            }
        }
    }

    // MARK: Rendezvous plumbing

    private func resumeProbe(_ result: Bool) {
        let cont: CheckedContinuation<Bool, Never>? = lock.withLock {
            let cont = probeContinuation
            probeContinuation = nil
            probeNonce = nil
            return cont
        }
        cont?.resume(returning: result)
    }

    private func resumeResync(_ value: UInt64?) {
        let cont: CheckedContinuation<UInt64?, Never>? = lock.withLock {
            let cont = resyncContinuation
            resyncContinuation = nil
            return cont
        }
        cont?.resume(returning: value)
    }

    // MARK: Direct-leg sequencing

    static func sequenced(seq: UInt64, frame: Data) -> Data {
        var out = Data(capacity: 8 + frame.count)
        out.append(contentsOf: withUnsafeBytes(of: seq.bigEndian, Array.init))
        out.append(frame)
        return out
    }

    static func desequence(_ raw: Data) -> (seq: UInt64, frame: Data)? {
        guard raw.count >= 8 else { return nil }
        return (raw.readUInt64(at: 0), Data(raw.dropFirst(8)))
    }
}

extension Data {
    /// Big-endian UInt64 at a relative offset (Data slices keep their parent's indices).
    fileprivate func readUInt64(at offset: Int) -> UInt64 {
        let start = startIndex + offset
        return (0..<8).reduce(UInt64(0)) { acc, i in
            (acc << 8) | UInt64(self[start + i])
        }
    }
}
