import ClientTLS
import Darwin
import Foundation
import WireProtocol

// MARK: - Typed transport errors

/// Transport-level errors with contract meaning for the engine (T-iOS-9).
///
/// Deliberately a SessionCore type, NOT WireProtocol: the frozen contract only
/// says "stream throw = transport error" (`TransportConnection.frames` doc);
/// WHICH error is a SessionCore concern — `SessionEngine` (T-iOS-10) lives in
/// the same module and classifies these.
///
/// Connect-time failures are NOT wrapped: `connect(to:)` rethrows the
/// underlying URLSession error verbatim, because the pairing probe
/// (`PairingError.classify`, T-iOS-8) inspects POSIX/NSURLError codes to map
/// `localNetworkDenied` / `atsBlocked` / `tlsFailure` — wrapping would break
/// that taxonomy.
public enum TermTransportError: Error, Equatable, Sendable {
    /// `receive()` failed with NSPOSIXErrorDomain EMSGSIZE(40) "Message too
    /// long" (ENOBUFS(55) accepted as fallback — T-iOS-2 spike measured; it is
    /// NOT a 1009 clean close on iOS): the single-frame ring-buffer replay
    /// exceeded `maximumMessageSize`. NON-RETRYABLE — the engine must surface
    /// `connection(.failed(.replayTooLarge))` and never feed this into the
    /// backoff loop (plan §1 / §3.2.1 coupling warning), otherwise the retry
    /// is deterministic and infinite.
    case replayTooLarge
    /// `send` on a connection that already terminated (client `close()`,
    /// server close frame, or transport error). Mirrors
    /// `FakeTransportError.sendAfterClose` (TestSupport).
    case sendAfterClose
}

// MARK: - Internal ping extension point (T-iOS-9 design decision)

/// INTERNAL, non-frozen ping hook.
///
/// Why this shape: the frozen `TransportConnection` (WireProtocol §3.1) is
/// exactly `{frames, send, close}` — no ping — and the frozen `TermTransport`
/// protocol must not change. But `PingScheduler` (T-iOS-5) needs a per-
/// connection `sendPing` the engine can call. Ruled-out alternatives:
/// a ping closure on `TransportConnection` (frozen-contract change → T-iOS-3),
/// and a transport-level registry keyed by connection (streams aren't
/// hashable; leak-prone). Chosen: transports that can ping conform to this
/// SessionCore-internal protocol; the engine downcasts its
/// `any TermTransport` and, when supported, opens the connection through
/// `connectPingable(to:)`, receiving the frozen `TransportConnection` PLUS a
/// per-connection `ConnectionPinger`. `FakeTransport` (TestSupport) does not
/// conform → engine tests drive `PingScheduler` through its injected closure
/// instead. No frozen contract is touched.
protocol PingableTermTransport: TermTransport {
    /// Like `connect(to:)`, but also hands back the ping hook for that same
    /// connection (for `PingScheduler.run(clock:sendPing:)` wiring).
    func connectPingable(to endpoint: HostEndpoint) async throws
        -> (connection: TransportConnection, pinger: any ConnectionPinger)
}

/// One live connection's ping capability (see `PingableTermTransport`).
protocol ConnectionPinger: Sendable {
    /// Send one WS ping; `true` iff the pong arrived in time (`false` = miss,
    /// counted by `PingScheduler` against `Tunables.pongMissLimit`).
    func ping() async -> Bool
}

// MARK: - Transport

/// Production `TermTransport` over `URLSessionWebSocketTask` (T-iOS-9).
///
/// Known platform pitfalls this type exists to fence (plan §1, spike-verified):
/// - Origin: stamped from `HostEndpoint.originHeader` ONLY (single source,
///   §5.1 — any string-assembled origin in this file is a review CRITICAL).
/// - `maximumMessageSize` defaults to 1 MiB; ring-buffer replay is one full
///   snapshot frame → every task gets `Tunables.maxWSMessageBytes` (16 MiB).
/// - `receive()` delivers ONE message per call → the loop re-arms forever.
/// - No automatic ping → `ConnectionPinger` (driven by `PingScheduler`).
public struct URLSessionTermTransport: TermTransport {
    /// `URLSessionWebSocketTask.maximumMessageSize` applied to every task.
    private let maxMessageBytes: Int
    /// C-iOS-2 (MEDIUM no-relaunch fix) · Resolve the optional device client
    /// identity LAZILY, once per `connect`, so a certificate imported mid-run
    /// takes effect on the NEXT connection without an app relaunch (a snapshot
    /// captured at composition would stay stale). When the resolved identity is
    /// present, the connection answers a `ClientCertificate` challenge (mTLS to a
    /// tunneled host) with its credential; when `nil`, the challenge is cancelled
    /// cleanly. Local (non-mTLS) hosts never issue the challenge, so a `nil`
    /// result is inert there.
    private let identityProvider: @Sendable () -> ClientIdentity?

    /// Fixed-identity convenience (snapshot callers / tests): wraps a constant
    /// provider, so behaviour is identical to capturing the identity directly.
    public init(identity: ClientIdentity? = nil) {
        self.init(identityProvider: { identity })
    }

    /// C-iOS-2 (MEDIUM no-relaunch fix) · per-connect identity resolution: the
    /// provider is re-consulted on every `connect`, so a nil→installed
    /// transition is picked up without relaunch.
    public init(identityProvider: @escaping @Sendable () -> ClientIdentity?) {
        self.init(
            maxMessageBytes: Tunables.maxWSMessageBytes, identityProvider: identityProvider
        )
    }

    /// Internal TEST seam: a shrunken cap makes the oversize→EMSGSIZE path
    /// deterministic without 16 MiB fixtures. Production code paths always go
    /// through `init()` / `init(identityProvider:)` and the frozen
    /// `Tunables.maxWSMessageBytes`.
    init(maxMessageBytes: Int, identity: ClientIdentity? = nil) {
        self.init(maxMessageBytes: maxMessageBytes, identityProvider: { identity })
    }

    init(
        maxMessageBytes: Int,
        identityProvider: @escaping @Sendable () -> ClientIdentity?
    ) {
        self.maxMessageBytes = maxMessageBytes
        self.identityProvider = identityProvider
    }

    public func connect(to endpoint: HostEndpoint) async throws -> TransportConnection {
        try await openConnection(to: endpoint).transportConnection()
    }

    /// Internal: concrete-typed connect (tests assert task configuration;
    /// `connectPingable` builds on it). The identity is resolved HERE, per
    /// connect, from `identityProvider` (no-relaunch pickup).
    func openConnection(to endpoint: HostEndpoint) async throws -> WSConnection {
        try await WSConnection.open(
            endpoint: endpoint, maxMessageBytes: maxMessageBytes,
            identity: identityProvider()
        )
    }
}

extension URLSessionTermTransport: PingableTermTransport {
    func connectPingable(to endpoint: HostEndpoint) async throws
        -> (connection: TransportConnection, pinger: any ConnectionPinger) {
        let connection = try await openConnection(to: endpoint)
        return (connection.transportConnection(), connection)
    }
}

// MARK: - Connection

/// One live WS connection: owns the mutable runtime handles (session, task,
/// receive loop) and hides them behind the frozen `TransportConnection`
/// capability struct.
///
/// Concurrency note: this is an NSLock-disciplined `@unchecked Sendable`
/// class, not an actor, because `URLSessionWebSocketDelegate` callbacks are
/// synchronous and ordering-sensitive (didOpen must resolve the handshake
/// continuation before didComplete can) — hopping them through actor `Task`s
/// would lose that ordering. Every mutable field is only touched under
/// `lock`; the URLSession handles themselves are thread-safe.
final class WSConnection: NSObject, @unchecked Sendable {
    private let lock = NSLock()
    private var handshakeContinuation: CheckedContinuation<Void, any Error>?
    private var isClosedByClient = false
    private var didObserveServerCloseFrame = false
    private var isTerminated = false
    private var droppedBinaryFrameCount = 0

    /// Set exactly once in `configure` (before any concurrent access) — IUO
    /// because `URLSession(configuration:delegate:…)` needs `self` post-init.
    private(set) var task: URLSessionWebSocketTask!
    private var session: URLSession!
    private var receiveLoop: Task<Void, Never>?
    private let frames: AsyncThrowingStream<String, any Error>
    private let framesContinuation: AsyncThrowingStream<String, any Error>.Continuation

    /// C-iOS-2 mTLS. Set once in `configure` BEFORE `task.resume()`, so the
    /// connection-level challenge (which can only arrive after resume) always
    /// reads a stable value without locking — same set-once discipline as
    /// `task`/`session`.
    private var identity: ClientIdentity?
    private let challengeResponder = MutualTLSChallengeResponder()

    private override init() {
        (frames, framesContinuation) = AsyncThrowingStream<String, any Error>.makeStream()
        super.init()
    }

    /// Connect: build session+task, resume, await the delegate-driven
    /// handshake (didOpen / didCompleteWithError — no receive-error guessing),
    /// then start the re-arming receive loop.
    static func open(
        endpoint: HostEndpoint, maxMessageBytes: Int, identity: ClientIdentity? = nil
    ) async throws -> WSConnection {
        let connection = WSConnection()
        connection.configure(
            endpoint: endpoint, maxMessageBytes: maxMessageBytes, identity: identity
        )
        try await connection.performHandshake()
        connection.startReceiveLoop()
        return connection
    }

    /// Wrap as the frozen boundary struct. Called exactly once per connection
    /// (`AsyncThrowingStream` is single-consumer).
    func transportConnection() -> TransportConnection {
        TransportConnection(
            frames: frames,
            send: { [self] in try await send($0) },
            close: { [self] in await close() }
        )
    }

    // MARK: Setup

    private func configure(
        endpoint: HostEndpoint, maxMessageBytes: Int, identity: ClientIdentity?
    ) {
        // Set BEFORE building the task/resume: the mTLS challenge can only fire
        // after `task.resume()`, so this write happens-before any read of it.
        self.identity = identity
        // Origin: SINGLE source of truth = endpoint.originHeader (plan §5.1).
        var request = URLRequest(url: endpoint.wsURL)
        request.setValue(endpoint.originHeader, forHTTPHeaderField: "Origin")
        let session = URLSession(configuration: .ephemeral, delegate: self, delegateQueue: nil)
        let task = session.webSocketTask(with: request)
        task.maximumMessageSize = maxMessageBytes
        self.session = session
        self.task = task
    }

    private func performHandshake() async throws {
        try await withCheckedThrowingContinuation {
            (continuation: CheckedContinuation<Void, any Error>) in
            lock.withLock { handshakeContinuation = continuation }
            task.resume()
        }
    }

    // MARK: Frozen-contract operations

    /// Send one client text frame; explicit `.sendAfterClose` once the
    /// connection terminated (client close, server close, or error).
    func send(_ frame: String) async throws {
        let isSendable = lock.withLock { !isClosedByClient && !isTerminated }
        guard isSendable else { throw TermTransportError.sendAfterClose }
        try await task.send(.string(frame))
    }

    /// Client detach: the server-side PTY keeps running (frozen-contract doc).
    /// Idempotent; finishes the frames stream cleanly.
    func close() async {
        lock.withLock {
            isClosedByClient = true
            isTerminated = true
        }
        task.cancel(with: .normalClosure, reason: nil)
        framesContinuation.finish()
        session.finishTasksAndInvalidate()
    }

    // MARK: Receive loop

    /// `URLSessionWebSocketTask.receive()` delivers exactly ONE message per
    /// call (plan §1 pitfall) — forgetting to re-arm silently stalls the
    /// stream, so this loops until the task terminates. The Task holds `self`
    /// strongly only while running; every termination path also invalidates
    /// the session (which drops its strong delegate reference to `self`).
    private func startReceiveLoop() {
        receiveLoop = Task { [self] in
            while true {
                do {
                    deliver(try await task.receive())
                } catch {
                    finishFrames(afterReceiveFailure: error)
                    return
                }
            }
        }
    }

    private func deliver(_ message: URLSessionWebSocketTask.Message) {
        switch message {
        case .string(let text):
            framesContinuation.yield(text)
        case .data:
            recordDroppedBinaryFrame()
        @unknown default:
            recordDroppedBinaryFrame()
        }
    }

    /// The server only ever sends text frames (src/protocol.ts), but it is an
    /// untrusted input source (standards §4): binary/unknown frames are
    /// dropped and counted — never fatal, receiving continues.
    private func recordDroppedBinaryFrame() {
        lock.withLock { droppedBinaryFrameCount += 1 }
    }

    /// Dropped-frame counter (test/diagnostic visibility).
    var droppedBinaryFrames: Int {
        lock.withLock { droppedBinaryFrameCount }
    }

    // MARK: Termination classification

    /// Classify how the connection ended (finish/throw are idempotent — the
    /// first termination path wins, later calls are no-ops):
    /// - clean (client `close()`, server close frame observed by the delegate,
    ///   or a processed close frame per `task.closeCode`) → stream FINISHES;
    /// - EMSGSIZE(40)/ENOBUFS(55) → typed `.replayTooLarge` (non-retryable);
    /// - anything else → stream THROWS the underlying error (retryable class).
    private func finishFrames(afterReceiveFailure error: any Error) {
        let wasClosedCleanly = lock.withLock {
            isClosedByClient || didObserveServerCloseFrame
        } || task.closeCode != .invalid
        markTerminated()
        if wasClosedCleanly {
            framesContinuation.finish()
        } else if Self.isMessageTooLong(error) {
            framesContinuation.finish(throwing: TermTransportError.replayTooLarge)
        } else {
            framesContinuation.finish(throwing: error)
        }
        session.finishTasksAndInvalidate()
    }

    private func markTerminated() {
        lock.withLock { isTerminated = true }
    }

    private func takeHandshakeContinuation() -> CheckedContinuation<Void, any Error>? {
        lock.withLock {
            defer { handshakeContinuation = nil }
            return handshakeContinuation
        }
    }

    /// T-iOS-2 spike (measured on a real server): exceeding
    /// `maximumMessageSize` fails `receive()` with NSPOSIXErrorDomain
    /// EMSGSIZE(40) "Message too long" — not a 1009 close. ENOBUFS(55) is the
    /// documented fallback for the same classification (plan §7 T-iOS-9
    /// errata). Checked at the top level and one underlying level.
    private static func isMessageTooLong(_ error: any Error) -> Bool {
        func matches(_ ns: NSError) -> Bool {
            ns.domain == NSPOSIXErrorDomain
                && (ns.code == Int(EMSGSIZE) || ns.code == Int(ENOBUFS))
        }
        let ns = error as NSError
        if matches(ns) { return true }
        guard let underlying = ns.userInfo[NSUnderlyingErrorKey] as? NSError else {
            return false
        }
        return matches(underlying)
    }
}

// MARK: - URLSession delegate (drives state — not receive-error guessing)

extension WSConnection: URLSessionWebSocketDelegate {
    func urlSession(
        _ session: URLSession,
        webSocketTask: URLSessionWebSocketTask,
        didOpenWithProtocol protocol: String?
    ) {
        takeHandshakeContinuation()?.resume()
    }

    /// C-iOS-2 · Connection-level auth challenge (server trust + client cert)
    /// arrives task-level for a WS task. Delegate the decision to the pure
    /// `MutualTLSChallengeResponder`: present the device identity for an mTLS
    /// tunnel host, default-handle the LE server trust, cancel cleanly if a
    /// client cert is demanded but none is installed. `identity` is set-once
    /// before `resume()` (see `configure`), so this read needs no lock.
    func urlSession(
        _ session: URLSession,
        task: URLSessionTask,
        didReceive challenge: URLAuthenticationChallenge,
        completionHandler: @escaping (URLSession.AuthChallengeDisposition, URLCredential?) -> Void
    ) {
        let resolution = challengeResponder.resolve(challenge, identity: identity)
        completionHandler(resolution.disposition, resolution.credential)
    }

    /// Server close frame processed → CLEAN close: the stream FINISHES
    /// (distinguishable from the transport-error THROW path).
    func urlSession(
        _ session: URLSession,
        webSocketTask: URLSessionWebSocketTask,
        didCloseWith closeCode: URLSessionWebSocketTask.CloseCode,
        reason: Data?
    ) {
        lock.withLock {
            didObserveServerCloseFrame = true
            isTerminated = true
        }
        framesContinuation.finish()
        session.finishTasksAndInvalidate()
    }

    func urlSession(
        _ session: URLSession,
        task: URLSessionTask,
        didCompleteWithError error: (any Error)?
    ) {
        if let handshake = takeHandshakeContinuation() {
            // Handshake never opened. Rethrow the UNDERLYING error verbatim —
            // PairingError.classify (T-iOS-8) depends on its POSIX/NSURLError
            // identity (see TermTransportError doc).
            handshake.resume(throwing: error ?? URLError(.cannotConnectToHost))
            session.finishTasksAndInvalidate()
            return
        }
        // Post-open completion: the receive loop normally observes the failure
        // first; this is the belt-and-braces path (idempotent finish).
        guard let error else {
            markTerminated()
            framesContinuation.finish()
            session.finishTasksAndInvalidate()
            return
        }
        finishFrames(afterReceiveFailure: error)
    }
}

// MARK: - Ping (internal hook implementation)

extension WSConnection: ConnectionPinger {
    /// One WS ping; `true` iff the pong arrived before the deadline.
    ///
    /// Deadline REUSES `Tunables.pingInterval` (no new constant — adding one
    /// would be a T-iOS-3 contract change): a pong that has not arrived by the
    /// time the NEXT ping would be due is by definition a miss;
    /// `PingScheduler` tolerates `Tunables.pongMissLimit - 1` consecutive
    /// misses before declaring the connection lost. Built on a first-value
    /// AsyncStream race (not a task group) so a black-holed
    /// `pongReceiveHandler` can never hang the caller.
    func ping() async -> Bool {
        let (verdicts, continuation) = AsyncStream<Bool>.makeStream()
        task.sendPing { error in continuation.yield(error == nil) }
        let deadline = Task {
            try? await Task.sleep(for: Tunables.pingInterval)
            continuation.yield(false)
        }
        defer {
            deadline.cancel()
            continuation.finish()
        }
        var iterator = verdicts.makeAsyncIterator()
        return await iterator.next() ?? false
    }
}
