import Foundation
import Network

/// A single-shot localhost HTTP listener that catches an OAuth redirect during a remote agent
/// sign-in (docs/REMOTE_AGENT_LOGIN.md §6) — the phone-side sibling of the Mac's `OAuthLoopback`.
///
/// The browser presenting the vendor's consent page runs ON THIS PHONE, so `localhost` in the
/// vendor's redirect URI resolves here: bind the port *before* asking the host to begin (Codex's
/// OAuth client only permits the fixed `localhost:1455`; Claude accepts any port), then forward
/// the captured `code`/`state` to the host over the E2EE channel. Answers exactly one request
/// with a tiny "return to Nucleic" page, then tears down. Binds 127.0.0.1 only.
final class OAuthRedirectListener: @unchecked Sendable {
    struct Callback: Sendable {
        let code: String
        let state: String?
    }

    /// The port actually bound (the requested one, or the OS-granted ephemeral port).
    let port: UInt16

    private let listener: NWListener
    private let queue = DispatchQueue(label: "xyz.blakeslee.nucleic.remote.oauth-redirect")
    private let lock = NSLock()
    private var continuation: CheckedContinuation<Callback?, Never>?
    private var pending: Callback?
    private var finished = false

    private init(listener: NWListener, port: UInt16) {
        self.listener = listener
        self.port = port
    }

    /// Bind a loopback port and begin listening. Pass `port` to require a specific one (Codex's
    /// fixed 1455); omit it for an ephemeral port (Claude). Returns nil when the OS won't grant
    /// it (something else holds 1455) — the caller then begins without a port and the host falls
    /// back to a paste capture where the provider supports one.
    static func start(port desiredPort: UInt16? = nil) -> OAuthRedirectListener? {
        let params = NWParameters.tcp
        // Loopback only — this port must never be reachable off-device.
        let endpointPort: NWEndpoint.Port =
            desiredPort.flatMap { NWEndpoint.Port(rawValue: $0) } ?? .any
        params.requiredLocalEndpoint = NWEndpoint.hostPort(host: "127.0.0.1", port: endpointPort)
        params.allowLocalEndpointReuse = true
        guard let listener = try? NWListener(using: params) else { return nil }

        let ready = DispatchSemaphore(value: 0)
        let portBox = Box<UInt16?>(nil)
        listener.stateUpdateHandler = { state in
            switch state {
            case .ready:
                portBox.value = listener.port?.rawValue
                ready.signal()
            case .failed, .cancelled:
                ready.signal()
            default:
                break
            }
        }
        let instanceBox = Box<OAuthRedirectListener?>(nil)
        listener.newConnectionHandler = { connection in
            instanceBox.value?.handle(connection)
        }
        listener.start(queue: DispatchQueue(
            label: "xyz.blakeslee.nucleic.remote.oauth-redirect.listen"))
        // A loopback bind is immediate; cap the wait so a wedged listener can't hang the flow.
        guard ready.wait(timeout: .now() + 2) == .success, let bound = portBox.value else {
            listener.cancel()
            return nil
        }
        let instance = OAuthRedirectListener(listener: listener, port: bound)
        instanceBox.value = instance
        return instance
    }

    /// Await the captured redirect; nil on timeout or `stop()`. Single-shot — a redirect that
    /// raced ahead of this call was buffered and returns immediately.
    func waitForCallback(timeout: TimeInterval) async -> Callback? {
        let buffered: Callback? = lock.withLock {
            if let pending { self.pending = nil; return pending }
            return nil
        }
        if let buffered { return buffered }
        return await withTaskGroup(of: Callback?.self) { group in
            group.addTask { [self] in
                await withCheckedContinuation { cont in
                    let deliverNow: Callback?? = lock.withLock {
                        if finished { return .some(nil) }
                        if let pending { self.pending = nil; return .some(pending) }
                        continuation = cont
                        return nil
                    }
                    if let deliverNow { cont.resume(returning: deliverNow) }
                }
            }
            group.addTask {
                try? await Task.sleep(for: .seconds(max(1, timeout)))
                return nil
            }
            let first = await group.next() ?? nil
            group.cancelAll()
            self.stop()
            return first
        }
    }

    /// Tear down; resumes a pending wait with nil. Idempotent.
    func stop() {
        let cont: CheckedContinuation<Callback?, Never>? = lock.withLock {
            finished = true
            let cont = continuation
            continuation = nil
            return cont
        }
        listener.cancel()
        cont?.resume(returning: nil)
    }

    // MARK: - One-request HTTP

    private func handle(_ connection: NWConnection) {
        connection.start(queue: queue)
        receiveRequest(connection, buffer: Data())
    }

    private func receiveRequest(_ connection: NWConnection, buffer: Data) {
        connection.receive(minimumIncompleteLength: 1, maximumLength: 16 * 1024) {
            [weak self] data, _, isComplete, error in
            guard let self else { connection.cancel(); return }
            var buffer = buffer
            if let data { buffer.append(data) }
            // The request line + headers end with CRLFCRLF; we only need the request line.
            if buffer.range(of: Data("\r\n\r\n".utf8)) != nil || isComplete || error != nil {
                self.respond(connection, requestHead: buffer)
            } else if buffer.count < 64 * 1024 {
                self.receiveRequest(connection, buffer: buffer)
            } else {
                connection.cancel()
            }
        }
    }

    private func respond(_ connection: NWConnection, requestHead: Data) {
        let callback = Self.parseCallback(requestHead: requestHead)
        let body = callback != nil
            ? "<html><body style=\"font-family:-apple-system\"><h3>Signed in.</h3>You can return to Nucleic.</body></html>"
            : "<html><body style=\"font-family:-apple-system\"><h3>Waiting for sign-in…</h3></body></html>"
        let response = "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\n"
            + "Content-Length: \(body.utf8.count)\r\nConnection: close\r\n\r\n" + body
        connection.send(content: Data(response.utf8), completion: .contentProcessed { _ in
            connection.cancel()
        })
        guard let callback else { return }
        let cont: CheckedContinuation<Callback?, Never>? = lock.withLock {
            guard !finished else { return nil }
            if let cont = continuation {
                continuation = nil
                return cont
            }
            pending = callback  // redirect raced ahead of waitForCallback — buffer it
            return nil
        }
        cont?.resume(returning: callback)
    }

    /// Pull `code` (+ `state`) out of the redirect's request line: `GET /path?code=…&state=… HTTP/1.1`.
    static func parseCallback(requestHead: Data) -> Callback? {
        guard let head = String(data: requestHead, encoding: .utf8),
              let requestLine = head.components(separatedBy: "\r\n").first
        else { return nil }
        let parts = requestLine.components(separatedBy: " ")
        guard parts.count >= 2,
              let components = URLComponents(string: "http://localhost\(parts[1])"),
              let code = components.queryItems?.first(where: { $0.name == "code" })?.value,
              !code.isEmpty
        else { return nil }
        let state = components.queryItems?.first(where: { $0.name == "state" })?.value
        return Callback(code: code, state: state)
    }

    private final class Box<T>: @unchecked Sendable {
        var value: T
        init(_ value: T) { self.value = value }
    }
}
