import Foundation

// The Nucleic Private Relay data path (SYNC_PROTOCOL §3.2, mesh P2): a WebSocket per device
// into the host's relay room (`cloud/nucleic-edge`), carrying the same length-prefixed Noise
// frames as every other transport. The relay only ever sees ciphertext — admission is a
// two-tier token (long-lived membership → 2-minute connection token, `relayToken.ts`), never
// a content credential. This file is the platform-agnostic client leg shared by the phone
// and future mac-peer dialers; the host's room listener builds on `RelayWebSocket` from
// NucleicCore (it additionally speaks the `RelayEnvelope` demux).

/// Relay REST + WebSocket endpoints. One base URL serves both (the Worker routes by path),
/// so a dev override — a single `*.workers.dev` origin — swaps everything at once.
public enum RelayAPI {
    /// The production relay every distributed build talks to.
    public static let defaultBaseURL = URL(string: "https://relay.nucleic.blakeslee.xyz")!

    /// Resolve a base URL: an explicit override (pairing payload / dev setting) or the default.
    public static func baseURL(_ override: String?) -> URL {
        override.flatMap(URL.init(string:)) ?? defaultBaseURL
    }

    /// The `/relay` WebSocket upgrade URL for a connection token (`?t=` — the Worker accepts
    /// the token as a query param because `URLSessionWebSocketTask` can't set custom headers
    /// portably across proxies).
    public static func webSocketURL(base: URL, connectionToken: String) -> URL {
        var components = URLComponents(url: base, resolvingAgainstBaseURL: false)!
        components.scheme = components.scheme == "http" ? "ws" : "wss"
        components.path = "/relay"
        components.queryItems = [URLQueryItem(name: "t", value: connectionToken)]
        return components.url!
    }

    public enum Error: Swift.Error, LocalizedError, Equatable {
        case membershipRejected(status: Int)
        case badResponse
        case socketClosed

        public var errorDescription: String? {
            switch self {
            case .membershipRejected(let status):
                status == 403
                    ? "This device's relay access was revoked — pair again to restore it."
                    : "The relay didn't accept this device's credential (\(status)) — reconnect over LAN once to refresh it."
            case .badResponse: "The relay returned an unexpected response."
            case .socketClosed: "The relay connection closed."
            }
        }
    }

    /// Trade a long-lived membership token for the short-lived connection token the WebSocket
    /// upgrade requires (`POST /v1/relay/connect`). Both tiers are HMAC-stateless server-side;
    /// this is one small authenticated POST right before dialing.
    public static func connectionToken(
        base: URL, membershipToken: String, session: URLSession = .shared
    ) async throws -> String {
        var request = URLRequest(url: base.appendingPathComponent("/v1/relay/connect"))
        request.httpMethod = "POST"
        request.timeoutInterval = 10
        request.setValue("Bearer \(membershipToken)", forHTTPHeaderField: "Authorization")
        let (data, response) = try await session.data(for: request)
        guard let http = response as? HTTPURLResponse else { throw Error.badResponse }
        guard http.statusCode == 200 else { throw Error.membershipRejected(status: http.statusCode) }
        guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
              let token = json["token"] as? String
        else { throw Error.badResponse }
        return token
    }
}

/// The room's presence broadcast (`room.ts broadcastPresence`) — the relay's text control
/// plane, explicitly non-secret. The host uses it to reap virtual channels for departed
/// devices (the room socket outlives any one peer, so nothing else signals a disconnect);
/// the phone uses it to fail fast when no host holds the room (its frames would be dropped).
public struct RelayPresence: Sendable, Equatable, Decodable {
    public struct Peer: Sendable, Equatable, Decodable {
        public let deviceId: String
        public let role: String

        public init(deviceId: String, role: String) {
            self.deviceId = deviceId
            self.role = role
        }
    }

    public let peers: [Peer]

    public init(peers: [Peer]) {
        self.peers = peers
    }

    public var hasHost: Bool { peers.contains { $0.role == "host" } }

    /// Parse a relay text frame; nil for anything but a presence broadcast.
    public static func parse(_ text: String) -> RelayPresence? {
        struct Envelope: Decodable {
            let t: String
            let peers: [Peer]?
        }
        guard let envelope = try? JSONDecoder().decode(Envelope.self, from: Data(text.utf8)),
              envelope.t == "presence"
        else { return nil }
        return RelayPresence(peers: envelope.peers ?? [])
    }
}

/// What the host's room listener needs from its socket — seam for tests (an in-memory fake
/// drives the demux without a network). `RelayWebSocket` is the real implementation.
public protocol RelayRoomSocket: AnyObject, Sendable {
    func messages() -> AsyncStream<RelayWebSocket.Message>
    func send(_ data: Data)
    func close()
}

/// A connected relay WebSocket: ordered sends, an inbound message stream, and a keepalive
/// ping loop. Shared by the phone's `RelayFrameChannel` (raw Noise frames) and the host's
/// room listener (enveloped frames + presence text). `connect` returns only after a
/// ping/pong round-trip confirms the upgrade was admitted.
public final class RelayWebSocket: RelayRoomSocket, @unchecked Sendable {
    public enum Message: Sendable {
        case binary(Data)
        case text(String)
    }

    private let task: URLSessionWebSocketTask
    private let inbound: AsyncStream<Message>
    private let inboundContinuation: AsyncStream<Message>.Continuation
    private let outboundContinuation: AsyncStream<Data>.Continuation
    private var writerTask: Task<Void, Never>?
    private var receiveTask: Task<Void, Never>?
    private var pingTask: Task<Void, Never>?

    /// How often to ping: keeps NAT bindings warm and detects a dead socket well inside the
    /// reconnect budget. (Cloudflare's edge answers protocol pings without waking the DO.)
    private static let pingInterval: Duration = .seconds(25)

    public static func connect(url: URL, session: URLSession = .shared) async throws -> RelayWebSocket {
        let task = session.webSocketTask(with: url)
        task.maximumMessageSize = WireFraming.maxFrameSize + 64
        task.resume()
        // A rejected upgrade (bad/expired token) surfaces as this ping failing.
        try await withCheckedThrowingContinuation { (c: CheckedContinuation<Void, Swift.Error>) in
            task.sendPing { error in
                if let error { c.resume(throwing: error) } else { c.resume() }
            }
        }
        return RelayWebSocket(task: task)
    }

    private init(task: URLSessionWebSocketTask) {
        self.task = task
        var inCont: AsyncStream<Message>.Continuation!
        inbound = AsyncStream(bufferingPolicy: .unbounded) { inCont = $0 }
        inboundContinuation = inCont

        var outCont: AsyncStream<Data>.Continuation!
        let outbound = AsyncStream<Data>(bufferingPolicy: .unbounded) { outCont = $0 }
        outboundContinuation = outCont

        // Single writer so wire order matches send order (the Noise nonce sequence).
        writerTask = Task { [task] in
            for await data in outbound {
                do { try await task.send(.data(data)) } catch { break }
            }
        }
        receiveTask = Task { [weak self, task] in
            while !Task.isCancelled {
                do {
                    switch try await task.receive() {
                    case .data(let data): self?.inboundContinuation.yield(.binary(data))
                    case .string(let text): self?.inboundContinuation.yield(.text(text))
                    @unknown default: break
                    }
                } catch {
                    break
                }
            }
            self?.inboundContinuation.finish()
        }
        pingTask = Task { [weak self, task] in
            while !Task.isCancelled {
                try? await Task.sleep(for: Self.pingInterval)
                guard !Task.isCancelled else { break }
                task.sendPing { [weak self] error in
                    if error != nil { self?.close() }
                }
            }
        }
    }

    /// Ordered inbound messages; finishes when the socket dies.
    public func messages() -> AsyncStream<Message> { inbound }

    /// Enqueue one binary message (ordered, non-blocking).
    public func send(_ data: Data) { outboundContinuation.yield(data) }

    /// Enqueue one text message (relay control plane — presence/register/wake, never E2EE
    /// payload). Ordering vs. binary sends is not guaranteed; control messages don't need it.
    public func sendText(_ text: String) {
        Task { [task] in try? await task.send(.string(text)) }
    }

    public func close() {
        pingTask?.cancel()
        receiveTask?.cancel()
        writerTask?.cancel()
        outboundContinuation.finish()
        inboundContinuation.finish()
        task.cancel(with: .goingAway, reason: nil)
    }
}

/// The phone-side relay `FrameChannel` (the client leg — raw Noise frames, no routing
/// envelope; the Room DO addresses clients itself). Frames keep the same `WireFraming`
/// length prefix as TCP so every layer above sees identical bytes on every transport.
public final class RelayFrameChannel: FrameChannel, @unchecked Sendable {
    public let peerDescription: String
    private let socket: RelayWebSocket
    private let accumulator = FrameAccumulator()
    private let stream: AsyncStream<Data>
    private let continuation: AsyncStream<Data>.Continuation
    private var pumpTask: Task<Void, Never>?

    /// Trade the membership token for a connection token and open the room WebSocket.
    public static func dial(
        base: URL, membershipToken: String, session: URLSession = .shared
    ) async throws -> RelayFrameChannel {
        let token = try await RelayAPI.connectionToken(
            base: base, membershipToken: membershipToken, session: session)
        let socket = try await RelayWebSocket.connect(
            url: RelayAPI.webSocketURL(base: base, connectionToken: token), session: session)
        return RelayFrameChannel(socket: socket, host: base.host ?? "relay")
    }

    init(socket: RelayWebSocket, host: String) {
        self.socket = socket
        self.peerDescription = "relay:\(host)"
        var cont: AsyncStream<Data>.Continuation!
        stream = AsyncStream(bufferingPolicy: .unbounded) { cont = $0 }
        continuation = cont
        pumpTask = Task { [weak self] in
            for await message in socket.messages() {
                guard let self else { break }
                switch message {
                case .binary(let data):
                    guard let frames = try? self.accumulator.push(data) else {
                        self.close()
                        return
                    }
                    for frame in frames { self.continuation.yield(frame) }
                case .text(let text):
                    // No host in the room means our frames are being dropped (`routeFromClient`
                    // has no host socket) — fail fast so the dial chain / retry loop moves on.
                    if let presence = RelayPresence.parse(text), !presence.hasHost {
                        self.close()
                        return
                    }
                }
            }
            self?.continuation.finish()
        }
    }

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

    public func send(_ frame: Data) { socket.send(WireFraming.frame(frame)) }

    public func close() {
        pumpTask?.cancel()
        socket.close()
        continuation.finish()
    }
}
