import Foundation
import IrdP2PCrypto

/// What the relay session reports to its owner.
public enum RelayEvent: Sendable, Equatable {
    /// The socket is open and the room joined; `roster` follows immediately.
    case joined
    /// Who else is in the room, as answered to our `join` (never includes us).
    case roster([PeerId])
    case peerJoined(PeerId)
    case peerLeft(PeerId)
    /// A forwarded payload; `from` is relay-stamped.
    case signal(from: PeerId, payload: JSONValue)
    /// The relay answered an `error` frame; the connection stays open.
    case relayError(RelayErrorCode)
    /// The connection dropped; a reconnect is scheduled after `after`.
    case reconnecting(attempt: Int, after: Duration)
    /// The session ended for good: `leave()` was called or the upgrade was refused permanently.
    case closed(reason: String)
}

/// One `/ws` connection to a room, with join, keepalive, reconnection (same peer
/// id, so the relay replaces the stale connection), pacing against the node's
/// budgets and frame-size checks. Owns nothing about payload semantics.
public actor RelaySession {
    public struct Configuration: Sendable {
        public var url: URL
        public var origin: String?
        public var room: RoomId
        public var peer: PeerId
        public var limits: AppDescriptor.Limits
        public var keepaliveInterval: Duration = .seconds(20)
        public var pongTimeout: Duration = .seconds(10)
        public var missedPongsBeforeReconnect = 2
        public var connectTimeout: TimeInterval = 15
        public var pacing = true
        public var jitter: @Sendable () -> Double = { Double.random(in: 0..<1) }

        public init(url: URL, origin: String?, room: RoomId, peer: PeerId, limits: AppDescriptor.Limits) {
            self.url = url
            self.origin = origin
            self.room = room
            self.peer = peer
            self.limits = limits
        }
    }

    public enum State: Sendable, Equatable {
        case idle, connecting, joined, reconnecting(attempt: Int), closed
    }

    public let configuration: Configuration
    private let transport: any WebSocketTransport
    private let clock: any P2PClock
    private var pacer: RelayPacer

    public private(set) var state: State = .idle
    public nonisolated let events: AsyncStream<RelayEvent>
    private let emit: AsyncStream<RelayEvent>.Continuation

    private var connection: (any WebSocketConnection)?
    private var runner: Task<Void, Never>?
    private var keepalive: Task<Void, Never>?
    private var reconnectSleep: Task<Void, Never>?
    private var attempt = 0
    private var startContinuation: CheckedContinuation<Void, any Error>?
    private var pongsAwaited = 0

    public init(configuration: Configuration, transport: any WebSocketTransport, clock: any P2PClock = SystemClock()) {
        self.configuration = configuration
        self.transport = transport
        self.clock = clock
        pacer = RelayPacer(limits: configuration.limits, now: clock.now)
        pacer.enabled = configuration.pacing
        (events, emit) = AsyncStream<RelayEvent>.makeStream(bufferingPolicy: .unbounded)
    }

    public var isJoined: Bool { state == .joined }

    /// Connects and joins. Throws when the first attempt fails permanently
    /// (`upgradeRejected`, `relay(.roomFull)`, `relay(.tooManyRooms)`); transient
    /// failures reconnect in the background instead.
    public func start() async throws {
        guard state == .idle else { return }
        try await withCheckedThrowingContinuation { (c: CheckedContinuation<Void, any Error>) in
            startContinuation = c
            runner = Task { await self.run() }
        }
    }

    /// Sends a `signal` frame, paced against the node's budgets.
    public func send(signal payload: JSONValue, to: PeerId) async throws {
        let text = ClientFrame.signal(to: to, payload: payload).encoded
        let size = text.utf8.count
        guard size <= configuration.limits.maxFrameBytes else {
            throw P2PError.frameTooLarge(bytes: size, limit: configuration.limits.maxFrameBytes)
        }
        let wait = pacer.take(byteCount: size, now: clock.now)
        if wait > .zero { try await clock.sleep(for: wait) }
        guard state == .joined, let connection else { throw P2PError.notJoined }
        try await connection.send(text: text)
    }

    /// The app came to the foreground or the network changed: reconnect now if not joined.
    public func wake() {
        attempt = 0
        if case .reconnecting = state { reconnectSleep?.cancel() }
    }

    /// Leaves the room and closes the socket; the session cannot be restarted.
    public func leave() async {
        guard state != .closed else { return }
        let conn = connection
        state = .closed
        runner?.cancel()
        keepalive?.cancel()
        reconnectSleep?.cancel()
        connection = nil
        if let conn {
            try? await conn.send(text: ClientFrame.leave.encoded)
            await conn.close(code: 1000, reason: "leave")
        }
        emit.yield(.closed(reason: "left"))
        emit.finish()
    }

    // MARK: - connection loop

    private func run() async {
        while !Task.isCancelled, state != .closed {
            state = attempt == 0 ? .connecting : .reconnecting(attempt: attempt)
            do {
                let conn = try await connect()
                connection = conn
                try await conn.send(text: ClientFrame.join(room: configuration.room, peer: configuration.peer).encoded)
                try await awaitJoin(conn)
                attempt = 0
                startKeepalive(conn)
                await receiveLoop(conn)
                stopKeepalive()
                connection = nil
                if state == .closed { return }
            } catch let error as P2PError {
                let dead = connection
                connection = nil
                stopKeepalive()
                if let dead { await dead.close(code: 1000, reason: nil) }
                if isTerminal(error) {
                    state = .closed
                    startContinuation?.resume(throwing: error)
                    startContinuation = nil
                    emit.yield(.closed(reason: "\(error)"))
                    emit.finish()
                    return
                }
                if let c = startContinuation, case .upgradeRejected = error {
                    // 429 on the very first attempt: still retry in the background, but tell the caller.
                    startContinuation = nil
                    c.resume()
                }
            } catch {
                connection = nil
                stopKeepalive()
                if state == .closed { return }
            }
            // reconnect
            let delay = Backoff.relay(attempt: attempt, jitter: configuration.jitter())
            attempt += 1
            state = .reconnecting(attempt: attempt)
            emit.yield(.reconnecting(attempt: attempt, after: delay))
            let sleeper = Task { _ = try? await clock.sleep(for: delay) }
            reconnectSleep = sleeper
            _ = await sleeper.value
            reconnectSleep = nil
        }
    }

    private func isTerminal(_ error: P2PError) -> Bool {
        switch error {
        case .upgradeRejected(let status): return status == 403 || status == 404
        case .relay(.roomFull), .relay(.tooManyRooms), .relay(.badMessage): return true
        case .closed: return state == .closed
        default: return false
        }
    }

    private func connect() async throws -> any WebSocketConnection {
        var request = URLRequest(url: configuration.url, timeoutInterval: configuration.connectTimeout)
        if let origin = configuration.origin { request.setValue(origin, forHTTPHeaderField: "Origin") }
        return try await transport.connect(request)
    }

    /// Reads frames until the `peers` roster (success) or a join error.
    private func awaitJoin(_ conn: any WebSocketConnection) async throws {
        while true {
            let frame = try await receiveFrame(conn)
            switch frame {
            case .peers(let roster):
                state = .joined
                startContinuation?.resume()
                startContinuation = nil
                emit.yield(.joined)
                emit.yield(.roster(roster))
                return
            case .error(let code):
                throw P2PError.relay(code)
            case .unknown, .pong:
                continue
            default:
                // Anything else before the roster is out of order; keep waiting.
                continue
            }
        }
    }

    private func receiveFrame(_ conn: any WebSocketConnection) async throws -> ServerFrame {
        while true {
            switch try await conn.receive() {
            case .text(let text):
                if let frame = try? ServerFrame.decode(text) { return frame }
                // Undecodable frames are ignored (forward compatibility).
            case .data:
                continue // binary frames are not part of the protocol
            }
        }
    }

    private func receiveLoop(_ conn: any WebSocketConnection) async {
        while state == .joined {
            do {
                let frame = try await receiveFrame(conn)
                switch frame {
                case .peers(let roster): emit.yield(.roster(roster))
                case .peerJoined(let p): emit.yield(.peerJoined(p))
                case .peerLeft(let p): emit.yield(.peerLeft(p))
                case .signal(let from, let payload): emit.yield(.signal(from: from, payload: payload))
                case .pong: pongsAwaited = 0
                case .error(let code):
                    if code == .rateLimited { pacer.drain(now: clock.now) }
                    emit.yield(.relayError(code))
                case .unknown: break
                }
            } catch {
                return
            }
        }
    }

    // MARK: - keepalive

    private func startKeepalive(_ conn: any WebSocketConnection) {
        pongsAwaited = 0
        keepalive = Task { [configuration] in
            while !Task.isCancelled {
                guard (try? await clock.sleep(for: configuration.keepaliveInterval)) != nil else { return }
                guard await self.pingOrDrop(conn) else { return }
            }
        }
    }

    private func pingOrDrop(_ conn: any WebSocketConnection) async -> Bool {
        guard state == .joined else { return false }
        if pongsAwaited >= configuration.missedPongsBeforeReconnect {
            // Half-open connection: drop it so the receive loop fails and reconnects.
            await conn.close(code: 1001, reason: "no pong")
            return false
        }
        pongsAwaited += 1
        do { try await conn.send(text: ClientFrame.ping.encoded) } catch { return false }
        return true
    }

    private func stopKeepalive() {
        keepalive?.cancel()
        keepalive = nil
    }
}

extension RelayPacer {
    /// After a `rate-limited` frame: empty both buckets so the next sends wait a full refill.
    mutating func drain(now: Duration) {
        _ = messages.take(messages.capacity, now: now)
        _ = bytes.take(bytes.capacity, now: now)
    }
}