import Foundation
import NucleicProtocol

/// The host's Nucleic Private 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) must succeed before the listener joins the composite;
/// 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 with capped backoff.
/// 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 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 stopped = false

    public init(dial: @escaping Dial) {
        self.dial = dial
    }

    /// 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() }
        guard let socket else {
            throw SyncTransportError.listenerFailed("relay listener started before connect()")
        }
        let (stream, continuation) = AsyncStream<any FrameChannel>.makeStream()
        self.continuation = continuation
        pump(socket)
        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
        lock.unlock()

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

    // 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 with capped backoff (re-minting tokens each attempt — connection tokens live
    /// two minutes).
    private func socketDied() {
        lock.lock()
        let wasStopped = stopped
        self.socket = nil
        let open = channels.values
        channels.removeAll()
        lock.unlock()
        for channel in open { channel.finish() }
        guard !wasStopped else { return }

        Task { [weak self] in
            var attempt = 0
            while let self {
                if self.locked({ self.stopped }) { return }
                attempt += 1
                let backoff = min(Double(attempt) * 2, 30)
                try? await Task.sleep(for: .seconds(backoff))
                do {
                    let socket = try await self.dial()
                    let stillWanted = self.locked {
                        if !self.stopped { self.socket = socket }
                        return !self.stopped
                    }
                    if !stillWanted {
                        socket.close()
                        return
                    }
                    self.pump(socket)
                    return
                } catch {
                    continue
                }
            }
        }
    }

    // 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).
    fileprivate func channelClosed(tag: Data) {
        lock.lock()
        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
    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)
        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() }
    }
}
