import Foundation
import NucleicProtocol

/// The host's Covalence Relay listener (SYNC_PROTOCOL §3.2, mesh P2): one outbound
/// WebSocket into this Mac's relay room, demultiplexed into a per-device virtual
/// `FrameChannel` per `RelayEnvelope` routing tag — so the `SyncHost` accept loop sees the
/// relay exactly like LAN or tailnet.
///
/// Wire shape on the host leg (only): each WebSocket message is
/// `[0x01][8-byte tag = sha256(deviceID)[..8]][WireFraming-framed Noise bytes]` — the Room DO
/// stamps the sender tag on client→host messages and routes host→client messages by the tag
/// we prepend. The client leg carries the raw framed bytes, so Noise and everything above see
/// identical bytes on every transport.
///
/// Lifecycle: `connect()` (async — mints the host membership, trades it for a connection
/// token, opens + verifies the socket) is attempted once up front so the common healthy case
/// commits with accurate status — but its failure no longer writes the method off until the
/// next restart: `start()` on an unconnected listener keeps redialing with capped backoff
/// (the first connect races LAN/tailnet bring-up and login-time network churn by design, and
/// losing that race is transient). A later socket drop closes every virtual channel (their
/// `ConnectionHandler`s wind down and clients see the room lose its host via presence) and
/// redials the same way. `statusChanged` reports each failed dial and the eventual
/// (re)connect so the Settings status row tracks reality instead of freezing on the first
/// error. Presence broadcasts reap channels for departed devices — the room socket outlives
/// any one peer, so nothing else would ever finish a dead channel's stream.
public final class RelayListener: SyncListener, @unchecked Sendable {
    public typealias Dial = @Sendable () async throws -> any RelayRoomSocket

    private let dial: Dial
    private let retryDelay: @Sendable (Int) -> Duration
    private let lock = NSLock()
    private var socket: (any RelayRoomSocket)?
    private var channels: [Data: RelayRoomChannel] = [:]
    private var continuation: AsyncStream<any FrameChannel>.Continuation?
    private var pumpTask: Task<Void, Never>?
    private var redialTask: Task<Void, Never>?
    private var watchdogTask: Task<Void, Never>?
    /// True while a redial loop is live — the single-flight guard (two concurrent loops
    /// would double-dial and double-report status).
    private var redialing = false
    /// Bumped on each (re)start of the redial loop. A `networkPathChanged` kick cancels the
    /// backing-off loop and starts a fresh one; this stamps the generations so the cancelled
    /// loop — which may still be mid-`dial()` — neither installs its late-arriving socket over
    /// the new loop's nor unlatches the `redialing` flag the new loop now owns.
    private var redialGeneration = 0
    private var stopped = false

    /// How often the watchdog checks that a disconnected listener is actually redialing.
    /// Deliberately coarse — it is a backstop for a leaked/died redial loop (seen in the
    /// field: Settings frozen on "Retrying…" with no dial traffic at all), not the retry
    /// cadence itself.
    private let watchdogInterval: Duration

    /// Live room-socket health: a failure description per failed (re)dial, nil on each
    /// (re)connect. Fires from the redial task — hop actors as needed. Set it before
    /// `connect()`/`start()`; it is not synchronized for later mutation.
    public var statusChanged: (@Sendable (String?) -> Void)?

    public init(
        dial: @escaping Dial,
        retryDelay: @escaping @Sendable (Int) -> Duration = { .seconds(min(Double($0) * 2, 30)) },
        watchdogInterval: Duration = .seconds(45)
    ) {
        self.dial = dial
        self.retryDelay = retryDelay
        self.watchdogInterval = watchdogInterval
    }

    /// Whether the room socket is currently up (the start commit point reads this to
    /// resolve races between a late first connect and the health-row snapshot).
    public var isConnected: Bool { locked { socket != nil } }

    /// Run `body` under the listener lock. Synchronous on purpose — `NSLock.lock` is
    /// unavailable from async contexts, so every async path funnels its state access here.
    private func locked<T>(_ body: () -> T) -> T {
        lock.lock()
        defer { lock.unlock() }
        return body()
    }

    /// Open (and verify) the first room socket. Called before the listener joins the
    /// composite so an unreachable/unenrollable relay degrades to a Settings status row
    /// instead of a silently-green method.
    public func connect() async throws {
        let socket = try await dial()
        locked { self.socket = socket }
    }

    public func start() throws -> AsyncStream<any FrameChannel> {
        lock.lock()
        defer { lock.unlock() }
        let (stream, continuation) = AsyncStream<any FrameChannel>.makeStream()
        self.continuation = continuation
        if let socket {
            pump(socket)
        } else {
            // connect() failed (or was never called): self-heal instead of starting dead. We already
            // hold `lock` here, so use the lock-held half — calling `redial()` (which re-acquires the
            // non-recursive `lock`) would reentrant-deadlock this thread and wedge the listener.
            redialLocked()
        }
        startWatchdogLocked()
        return stream
    }

    public func stop() {
        lock.lock()
        stopped = true
        let socket = self.socket
        self.socket = nil
        let open = channels.values
        channels.removeAll()
        let continuation = self.continuation
        self.continuation = nil
        pumpTask?.cancel()
        pumpTask = nil
        redialTask?.cancel()
        redialTask = nil
        watchdogTask?.cancel()
        watchdogTask = nil
        lock.unlock()

        socket?.close()
        for channel in open { channel.finish() }
        continuation?.finish()
    }

    /// The backstop for a redial loop that died without reconnecting (whatever the path —
    /// every future leak included): while the listener is started and the room socket is
    /// down, make sure a redial loop exists. Without this, one lost loop froze the Settings
    /// row on "Retrying…" forever with zero dial traffic behind it.
    private func startWatchdogLocked() {
        guard watchdogTask == nil else { return }
        let interval = watchdogInterval
        watchdogTask = Task { [weak self] in
            while !Task.isCancelled {
                try? await Task.sleep(for: interval)
                guard !Task.isCancelled, let self else { return }
                let needsRedial = self.locked {
                    !self.stopped && self.socket == nil && !self.redialing
                }
                if needsRedial { self.redial() }
            }
        }
    }

    // MARK: - Room pump (demux + presence + redial)

    private func pump(_ socket: any RelayRoomSocket) {
        pumpTask = Task { [weak self] in
            for await message in socket.messages() {
                guard let self, !Task.isCancelled else { return }
                switch message {
                case .binary(let data): self.route(data)
                case .text(let text):
                    if let presence = RelayPresence.parse(text) { self.reap(keeping: presence) }
                }
            }
            self?.socketDied()
        }
    }

    /// client → host: unwrap the envelope and feed the tagged device's virtual channel,
    /// creating (and yielding) it on first sight. Malformed envelopes are dropped, mirroring
    /// the DO.
    private func route(_ data: Data) {
        guard let (tag, payload) = RelayEnvelope.unwrap(data) else { return }
        lock.lock()
        var created: RelayRoomChannel?
        var channel = channels[tag]
        if channel == nil, !stopped {
            let fresh = RelayRoomChannel(tag: tag, listener: self)
            channels[tag] = fresh
            channel = fresh
            created = fresh
        }
        let continuation = self.continuation
        lock.unlock()

        if let created { continuation?.yield(created) }
        channel?.receive(payload)
    }

    /// Presence says who's in the room; any channel whose device left gets its stream
    /// finished so its `ConnectionHandler` winds down. (A device that *reconnects* keeps its
    /// tag — the DO supersedes its old socket silently — and converges through the fresh
    /// handshake failing the old handler.)
    private func reap(keeping presence: RelayPresence) {
        let liveTags = Set(presence.peers.filter { $0.role != "host" }
            .map { RelayEnvelope.tag(forDeviceID: $0.deviceId) })
        lock.lock()
        let dead = channels.filter { !liveTags.contains($0.key) }
        for key in dead.keys { channels.removeValue(forKey: key) }
        lock.unlock()
        for channel in dead.values { channel.finish() }
    }

    /// The room socket dropped: every relayed connection is gone. Close their channels and
    /// redial.
    private func socketDied() {
        lock.lock()
        let wasStopped = stopped
        let dead = self.socket
        self.socket = nil
        let open = channels.values
        channels.removeAll()
        lock.unlock()
        // Close the dead socket explicitly — its stream already finished, but its ping/writer
        // tasks (and the URLSession task they retain) only wind down on close().
        dead?.close()
        for channel in open { channel.finish() }
        guard !wasStopped else { return }
        statusChanged?(RelayAPI.Error.socketClosed.errorDescription ?? "The relay connection closed.")
        redial()
    }

    /// Dial until a fresh room socket is up (capped backoff, re-minting tokens each attempt
    /// — connection tokens live two minutes), then pump it. Shared by the launch path (the
    /// first connect lost its race with LAN/tailnet bring-up or network churn), a room
    /// socket dying mid-run, and the watchdog. Single-flight: a second call while a loop is
    /// live is a no-op.
    private func redial() {
        lock.lock()
        defer { lock.unlock() }
        redialLocked()
    }

    /// The lock-held half of `redial()` — for callers (`start`) already inside the lock.
    /// `immediate` (a `networkPathChanged` kick) skips the initial backoff so the first dial
    /// fires now instead of sitting out the redial cadence.
    private func redialLocked(immediate: Bool = false) {
        guard !stopped, !redialing else { return }
        redialing = true
        redialGeneration &+= 1
        let generation = redialGeneration
        redialTask = Task { [weak self] in
            // Clear the single-flight flag only if we're still the current loop: a
            // `networkPathChanged` kick supersedes us with a fresh loop that already owns
            // `redialing`, so our (cancelled) exit must not unlatch it.
            defer { self?.locked { if self?.redialGeneration == generation { self?.redialing = false } } }
            var attempt = 0
            var settleFirst = !immediate
            while let self, !Task.isCancelled {
                if self.locked({ self.stopped }) { return }
                // Back off before each dial — except the first dial of a path-change kick,
                // which is the whole point: don't sit out the remaining backoff (up to 30s).
                if settleFirst {
                    attempt += 1
                    try? await Task.sleep(for: self.retryDelay(attempt))
                }
                settleFirst = true
                if Task.isCancelled { return }
                do {
                    let socket = try await self.dial()
                    let stillWanted = self.locked {
                        // A superseding kick (newer generation) or a stop raced our dial —
                        // drop the socket we just opened rather than install it over theirs.
                        guard !self.stopped, self.redialGeneration == generation else { return false }
                        self.socket = socket
                        return true
                    }
                    if !stillWanted {
                        socket.close()
                        return
                    }
                    self.pump(socket)
                    self.statusChanged?(nil)
                    return
                } catch {
                    if Task.isCancelled { return }
                    self.statusChanged?(error.localizedDescription)
                    continue
                }
            }
        }
    }

    /// A network path change (Wi-Fi ⇄ wired ⇄ cellular, or offline → online) just fired. If the
    /// room socket is down, dial *now* rather than sitting out the redial loop's backoff (up to
    /// 30s) or the watchdog interval (45s) — that idle window is exactly when a cellular/WAN
    /// device finds the room hostless and can't reach this Mac. A live socket is left untouched:
    /// the `RelayWebSocket` keepalive ping reaps a silently-orphaned one within an interval, and
    /// recycling a healthy socket would needlessly tear down every relayed device channel.
    public func networkPathChanged() {
        lock.lock()
        defer { lock.unlock() }
        guard !stopped, socket == nil else { return }
        // Supersede any backing-off loop with one that dials immediately. Cancel it, drop the
        // single-flight latch, then restart: `redialLocked` bumps the generation, so the
        // cancelled loop can neither install a late socket nor unlatch the new loop's `redialing`.
        redialTask?.cancel()
        redialing = false
        redialLocked(immediate: true)
    }

    // MARK: - Channel callbacks

    /// host → client: envelope the framed bytes with the target tag; the DO strips the
    /// envelope and forwards the raw payload to exactly that device.
    fileprivate func sendToPeer(tag: Data, framed: Data) {
        lock.lock()
        let socket = self.socket
        lock.unlock()
        socket?.send(RelayEnvelope.wrap(framed, tag: tag))
    }

    /// A `ConnectionHandler` closed its channel — drop the demux slot so a future frame from
    /// that tag opens a fresh channel (and a fresh handshake). Identity-checked: generations
    /// of the same device share one tag, so a *superseded* handler's deferred close (the
    /// `SyncHost` dedup grace, the liveness reaper) must not evict its successor's live slot.
    fileprivate func channelClosed(tag: Data, channel: RelayRoomChannel) {
        lock.lock()
        if channels[tag] === channel { channels.removeValue(forKey: tag) }
        lock.unlock()
    }
}

/// One relayed device's `FrameChannel`: frames in from the listener's demux, frames out
/// through the listener's enveloped send. Closing it only forgets the demux slot — the room
/// socket is shared and stays up.
final class RelayRoomChannel: FrameChannel, @unchecked Sendable {
    let peerDescription: String
    let transport: SyncTransportHint? = .relay
    private let tag: Data
    private weak var listener: RelayListener?
    private let accumulator = FrameAccumulator()
    private let stream: AsyncStream<Data>
    private let continuation: AsyncStream<Data>.Continuation
    private let lock = NSLock()
    private var finished = false

    init(tag: Data, listener: RelayListener) {
        self.tag = tag
        self.listener = listener
        self.peerDescription = "relay:" + tag.map { String(format: "%02x", $0) }.joined()
        var cont: AsyncStream<Data>.Continuation!
        stream = AsyncStream(bufferingPolicy: .unbounded) { cont = $0 }
        continuation = cont
    }

    func frames() -> AsyncStream<Data> { stream }

    func send(_ frame: Data) {
        listener?.sendToPeer(tag: tag, framed: WireFraming.frame(frame))
    }

    func close() {
        listener?.channelClosed(tag: tag, channel: self)
        finish()
    }

    /// Demuxed bytes for this device from the room socket (already de-enveloped).
    func receive(_ payload: Data) {
        guard let frames = try? accumulator.push(payload) else {
            // Oversize/garbage — drop the virtual connection, mirroring `LANChannel`.
            close()
            return
        }
        for frame in frames { continuation.yield(frame) }
    }

    func finish() {
        lock.lock()
        let wasFinished = finished
        finished = true
        lock.unlock()
        if !wasFinished { continuation.finish() }
    }
}
