import Foundation
import NucleicProtocol

// Remote agent sign-in, host side (docs/REMOTE_AGENT_LOGIN.md §4). The brokers own everything
// secret (PKCE, exchange, persistence); what varies is only how the consent URL reaches a
// browser and how the redirect comes back. `RemoteLoginChallenge`/`RemoteLoginCapture` are that
// seam: the local flow satisfies it with NSWorkspace + OAuthLoopback, the remote flow — this
// coordinator — with a wire challenge to the initiating device and a suspended continuation its
// `agentLoginCallback` resumes.

/// What a sign-in attempt needs the (remote) user's browser for: the consent URL, and how the
/// redirect must be captured. The broker builds it; the driver presents it.
public struct RemoteLoginChallenge: Sendable, Equatable {
    public enum Capture: Sendable, Equatable {
        /// The device listens on 127.0.0.1:`port` for `GET path?code=…&state=…` — the browser
        /// runs on the device, so `localhost` in the vendor redirect resolves there.
        case loopback(port: UInt16, path: String)
        /// The consent page renders `code#state` for the user to copy-paste (Claude's console
        /// redirect fallback).
        case pasteCode
    }

    public let authorizeURL: URL
    public let capture: Capture

    public init(authorizeURL: URL, capture: Capture) {
        self.authorizeURL = authorizeURL
        self.capture = capture
    }
}

/// A captured redirect: the authorization code, plus the redirect's `state` when the capture
/// carried one (loopback). A pasted `code#state` blob rides whole in `code`, state nil — the
/// exchange parses it, exactly as the local paste fallback does.
public struct RemoteLoginCode: Sendable, Equatable {
    public let code: String
    public let state: String?

    public init(code: String, state: String? = nil) {
        self.code = code
        self.state = state
    }
}

/// Present one challenge to the user's browser and return the captured redirect. Suspends for
/// the whole browser round-trip; throws to abort (device cancelled, connection lost, deadline).
public typealias RemoteLoginCapture = @Sendable (RemoteLoginChallenge) async throws -> RemoteLoginCode

/// Race a capture against the broker's login deadline. Nil on timeout, cancellation, or
/// capture failure — the caller folds nil into its provider-specific failure message. The
/// capture closure must resume on cancellation (`AgentLoginCoordinator.present` does, via its
/// cancellation handler) or the group would wait on it forever.
func remoteCaptureWithDeadline(
    _ challenge: RemoteLoginChallenge,
    capture: @escaping RemoteLoginCapture,
    timeoutNanos: UInt64
) async -> RemoteLoginCode? {
    try? await withThrowingTaskGroup(of: RemoteLoginCode.self) { group in
        group.addTask { try await capture(challenge) }
        group.addTask {
            try await Task.sleep(nanoseconds: timeoutNanos)
            throw CancellationError()
        }
        defer { group.cancelAll() }
        return try await group.next()!
    }
}

/// Runs remote sign-in attempts for a host: one flow per provider at a time, correlated by the
/// device-minted `requestID`. `run` drives a broker's `loginRemotely` end to end — pushing the
/// challenge at the initiating connection and suspending until its `agentLoginCallback` /
/// `agentLoginCancel` (or the broker deadline) resumes the capture — then answers with the
/// definitive `agentLoginResult`. Actor-isolated: the pending-capture table is the only state.
public actor AgentLoginCoordinator {
    /// Pushes one host message at the initiating connection (dropped if it closed).
    public typealias HostSend = @Sendable (HostMsg) async -> Void

    private let claude: ClaudeCredentialBroker?
    private let codex: CodexCredentialBroker?
    /// How long a pushed challenge stays answerable — mirrored to the device as `expiresAt`.
    /// The brokers' own login timeout is the enforcement; this is the advertised value.
    private let challengeTTL: TimeInterval
    private let now: @Sendable () -> Date

    private var pendingCaptures: [String: CheckedContinuation<RemoteLoginCode, Error>] = [:]
    private var providersInFlight: Set<AgentLoginProvider> = []

    public init(
        claude: ClaudeCredentialBroker?, codex: CodexCredentialBroker?,
        challengeTTL: TimeInterval = 10 * 60,
        now: @escaping @Sendable () -> Date = { Date() }
    ) {
        self.claude = claude
        self.codex = codex
        self.challengeTTL = challengeTTL
        self.now = now
    }

    /// Which providers this coordinator can broker — feeds
    /// `WireProviderAuthStatus.canBrokerLogin` and the host's `canBrokerAgentLogin` bit.
    public nonisolated var brokerableProviders: [AgentLoginProvider] {
        var providers: [AgentLoginProvider] = []
        if claude != nil { providers.append(.claude) }
        if codex != nil { providers.append(.codex) }
        return providers
    }

    /// Drive one attempt end to end and send its result; returns whether it succeeded (so the
    /// host can rebroadcast auth status / note the audit line). Long-running — call from a task,
    /// never from a connection's dispatch loop.
    public func run(_ request: WireAgentLoginBegin, send: @escaping HostSend) async -> Bool {
        let provider = request.provider
        // One flow per provider: the brokers FIFO-gate anyway (the loopback port and the
        // refresh lineage are both single-flight), but failing fast beats queueing a consent
        // page behind another user's minutes-long attempt.
        guard !providersInFlight.contains(provider) else {
            await send(.agentLoginResult(WireAgentLoginResult(
                requestID: request.requestID,
                error: "Another sign-in for this provider is already in progress on this host.")))
            return false
        }
        providersInFlight.insert(provider)
        defer {
            providersInFlight.remove(provider)
            discardPending(requestID: request.requestID)
        }

        let requestID = request.requestID
        let capture: RemoteLoginCapture = { [weak self] challenge in
            guard let self else { throw CancellationError() }
            return try await self.present(challenge, requestID: requestID, send: send)
        }

        let failure: String?
        switch provider {
        case .claude:
            guard let claude else {
                failure = "Claude sign-in is unavailable on this host."
                break
            }
            switch await claude.loginRemotely(
                boundLoopbackPort: request.boundLoopbackPort, capture: capture)
            {
            case .connected: failure = nil
            case .failed(let message): failure = message
            }
        case .codex:
            guard let codex else {
                failure = "Codex sign-in is unavailable on this host."
                break
            }
            switch await codex.loginRemotely(
                boundLoopbackPort: request.boundLoopbackPort, capture: capture)
            {
            case .connected: failure = nil
            case .failed(let message): failure = message
            }
        default:
            failure = "This host can't sign in to “\(provider.rawValue)”."
        }

        await send(.agentLoginResult(WireAgentLoginResult(requestID: requestID, error: failure)))
        return failure == nil
    }

    /// The initiating device answered with the captured redirect — resume the suspended flow.
    /// Stale/unknown ids are dropped (the attempt already timed out or was cancelled).
    public func deliverCallback(_ callback: WireAgentLoginCallback) {
        guard let pending = pendingCaptures.removeValue(forKey: callback.requestID) else { return }
        pending.resume(returning: RemoteLoginCode(code: callback.code, state: callback.state))
    }

    /// The device abandoned the attempt — fail the suspended flow (it answers with a failed
    /// result through its normal path).
    public func cancel(requestID: String) {
        discardPending(requestID: requestID)
    }

    // MARK: - Capture plumbing

    /// Push the challenge at the device, then suspend until `deliverCallback`/`cancel` (or task
    /// cancellation — the broker's deadline — via the cancellation handler) resumes it.
    private func present(
        _ challenge: RemoteLoginChallenge, requestID: String, send: @escaping HostSend
    ) async throws -> RemoteLoginCode {
        let wireCapture: AgentLoginCapture
        switch challenge.capture {
        case .loopback(let port, let path): wireCapture = .loopback(port: port, path: path)
        case .pasteCode: wireCapture = .pasteCode
        }
        await send(.agentLoginChallenge(WireAgentLoginChallenge(
            requestID: requestID,
            authorizeURL: challenge.authorizeURL.absoluteString,
            capture: wireCapture,
            expiresAt: now().addingTimeInterval(challengeTTL))))
        return try await withTaskCancellationHandler {
            try await withCheckedThrowingContinuation { continuation in
                // A second challenge for the same id can't happen (one flow per provider, one
                // capture per flow) — but resume-don't-leak if it ever does.
                if let stale = pendingCaptures.removeValue(forKey: requestID) {
                    stale.resume(throwing: CancellationError())
                }
                pendingCaptures[requestID] = continuation
            }
        } onCancel: {
            Task { await self.cancel(requestID: requestID) }
        }
    }

    private func discardPending(requestID: String) {
        guard let pending = pendingCaptures.removeValue(forKey: requestID) else { return }
        pending.resume(throwing: CancellationError())
    }
}
