import Foundation

#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
#if canImport(CryptoKit)
import CryptoKit
#else
import Crypto
#endif

/// The OAuth 2.0 + PKCE mechanics Nucleic uses to sign a user into their ChatGPT/Codex subscription
/// **itself**, without shelling out to `codex login` or requiring an external terminal — the Codex
/// counterpart to ``ClaudeOAuth``.
///
/// Codex's login is a public-client Authorization-Code-with-PKCE flow against
/// `auth.openai.com/oauth/authorize` (consent) and `auth.openai.com/oauth/token` (exchange +
/// refresh), whose client only permits the fixed redirect `http://localhost:1455/auth/callback`.
/// Nucleic mints the PKCE pair, opens consent in the browser, catches the redirect on that loopback,
/// exchanges the code, and writes the exact `~/.codex/auth.json` shape ``CodexAuthFile`` reconciles
/// — which the credential mesh then seals to runners and paired devices over Covalence. Token
/// requests are `application/x-www-form-urlencoded` (OpenAI's endpoint), unlike Claude's JSON.
///
/// The parsing/URL-building surface is pure and unit-tested; the two networked calls run through an
/// injectable ``Transport`` so tests exercise the flow without touching OpenAI.
public struct CodexOAuth: Sendable {
    // MARK: - Public client constants

    /// Codex CLI's public OAuth client id (a public client — PKCE is the proof, no secret).
    public static let clientID = "app_EMoamEEZ73f0CkXaXp7hrann"
    public static let authorizeEndpoint = "https://auth.openai.com/oauth/authorize"
    public static let tokenEndpoint = "https://auth.openai.com/oauth/token"
    /// The only redirect the Codex client allows: a fixed loopback port + path.
    public static let redirectURI = "http://localhost:1455/auth/callback"
    public static let loopbackPort: UInt16 = 1455
    public static let loopbackPath = "/auth/callback"
    public static let scopes = ["openid", "profile", "email", "offline_access"]

    // MARK: - Dependencies

    /// A single HTTP round-trip: request in, `(body, response)` out. Injectable so tests never hit
    /// the network; the default drives `URLSession`.
    public typealias Transport = @Sendable (URLRequest) async throws -> (Data, HTTPURLResponse)

    private let transport: Transport
    private let now: @Sendable () -> Date

    public init(
        transport: @escaping Transport = CodexOAuth.liveTransport,
        now: @escaping @Sendable () -> Date = { Date() }
    ) {
        self.transport = transport
        self.now = now
    }

    public enum OAuthError: Error, CustomStringConvertible, Equatable {
        case network(String)
        case httpStatus(Int, String)
        case malformedResponse
        case stateMismatch

        public var description: String {
            switch self {
            case .network(let m): return "Network error: \(m)"
            case .httpStatus(let code, _): return "Sign-in server returned HTTP \(code)"
            case .malformedResponse: return "Sign-in server returned an unexpected response"
            case .stateMismatch: return "Sign-in response did not match this request"
            }
        }
    }

    // MARK: - PKCE

    /// A one-shot PKCE + anti-forgery bundle for a single login attempt.
    public struct PKCE: Sendable, Equatable {
        public let verifier: String
        public let challenge: String
        public let state: String

        public static func generate() -> PKCE {
            let verifier = CodexOAuth.base64URL(randomBytes(32))
            let state = CodexOAuth.base64URL(randomBytes(32))
            return PKCE(verifier: verifier, challenge: challenge(for: verifier), state: state)
        }

        /// Deterministic constructor for tests.
        public init(verifier: String, state: String) {
            self.verifier = verifier
            self.challenge = CodexOAuth.PKCE.challenge(for: verifier)
            self.state = state
        }

        init(verifier: String, challenge: String, state: String) {
            self.verifier = verifier
            self.challenge = challenge
            self.state = state
        }

        static func challenge(for verifier: String) -> String {
            let digest = SHA256.hash(data: Data(verifier.utf8))
            return CodexOAuth.base64URL(Data(digest))
        }

        private static func randomBytes(_ count: Int) -> Data {
            var g = SystemRandomNumberGenerator()
            var bytes = Data(count: count)
            for i in 0..<count { bytes[i] = UInt8.random(in: 0...255, using: &g) }
            return bytes
        }
    }

    // MARK: - Authorize URL

    /// The consent URL to open in the browser. The extra `id_token_add_organizations` /
    /// `codex_cli_simplified_flow` / `originator` parameters match what the Codex CLI sends so the
    /// returned `id_token` carries the account claims Codex expects.
    public static func authorizeURL(pkce: PKCE) -> URL {
        var components = URLComponents(string: authorizeEndpoint)!
        components.queryItems = [
            URLQueryItem(name: "response_type", value: "code"),
            URLQueryItem(name: "client_id", value: clientID),
            URLQueryItem(name: "redirect_uri", value: redirectURI),
            URLQueryItem(name: "scope", value: scopes.joined(separator: " ")),
            URLQueryItem(name: "code_challenge", value: pkce.challenge),
            URLQueryItem(name: "code_challenge_method", value: "S256"),
            URLQueryItem(name: "id_token_add_organizations", value: "true"),
            URLQueryItem(name: "codex_cli_simplified_flow", value: "true"),
            URLQueryItem(name: "originator", value: "codex_cli_rs"),
            URLQueryItem(name: "state", value: pkce.state),
        ]
        return components.url!
    }

    // MARK: - Code exchange / refresh

    /// Exchange an authorization `code` for a credential, returning the `~/.codex/auth.json` JSON.
    public func exchange(code: String, pkce: PKCE) async throws -> String {
        let form = [
            "grant_type": "authorization_code",
            "code": code.trimmingCharacters(in: .whitespacesAndNewlines),
            "redirect_uri": Self.redirectURI,
            "client_id": Self.clientID,
            "code_verifier": pkce.verifier,
        ]
        let data = try await post(form)
        guard let json = Self.authFileJSON(fromTokenResponse: data, now: now()) else {
            throw OAuthError.malformedResponse
        }
        return json
    }

    /// Rotate a `refreshToken` into a fresh `auth.json`. Returns nil when the server
    /// rejects/garbles the exchange (the caller treats nil as "re-login required").
    public func refresh(refreshToken: String) async -> String? {
        let form = [
            "grant_type": "refresh_token",
            "refresh_token": refreshToken,
            "client_id": Self.clientID,
            "scope": Self.scopes.joined(separator: " "),
        ]
        guard let data = try? await post(form) else { return nil }
        return Self.authFileJSON(
            fromTokenResponse: data, now: now(), fallbackRefreshToken: refreshToken)
    }

    private func post(_ form: [String: String]) async throws -> Data {
        var request = URLRequest(url: URL(string: Self.tokenEndpoint)!)
        request.httpMethod = "POST"
        request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type")
        request.setValue("application/json", forHTTPHeaderField: "Accept")
        request.httpBody = Data(Self.formURLEncode(form).utf8)
        let (data, response): (Data, HTTPURLResponse)
        do {
            (data, response) = try await transport(request)
        } catch {
            throw OAuthError.network(String(describing: error))
        }
        guard (200..<300).contains(response.statusCode) else {
            throw OAuthError.httpStatus(response.statusCode, String(decoding: data, as: UTF8.self))
        }
        return data
    }

    // MARK: - Response → auth.json shape

    /// Build the persisted `auth.json` from a token-endpoint response. The response carries
    /// `id_token`, `access_token`, `refresh_token`; Codex derives the ChatGPT `account_id` from the
    /// id_token's `https://api.openai.com/auth` claim and stamps `last_refresh` (which
    /// ``CodexAuthFile`` uses for newest-wins). A refresh response may omit `refresh_token` — the
    /// lineage is unchanged, so `fallbackRefreshToken` preserves it.
    static func authFileJSON(
        fromTokenResponse data: Data,
        now: Date,
        fallbackRefreshToken: String? = nil
    ) -> String? {
        guard let obj = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
              let accessToken = obj["access_token"] as? String, !accessToken.isEmpty,
              let idToken = obj["id_token"] as? String, !idToken.isEmpty
        else { return nil }

        let refreshToken = (obj["refresh_token"] as? String).flatMap { $0.isEmpty ? nil : $0 }
            ?? fallbackRefreshToken
        guard let refreshToken, !refreshToken.isEmpty else { return nil }

        var tokens: [String: Any] = [
            "id_token": idToken,
            "access_token": accessToken,
            "refresh_token": refreshToken,
        ]
        if let accountID = accountID(fromIDToken: idToken) {
            tokens["account_id"] = accountID
        }

        let auth: [String: Any] = [
            "OPENAI_API_KEY": NSNull(),
            "tokens": tokens,
            "last_refresh": iso8601(now),
        ]
        guard let out = try? JSONSerialization.data(withJSONObject: auth),
              let string = String(data: out, encoding: .utf8)
        else { return nil }
        return string
    }

    /// The ChatGPT account id from an id_token's `https://api.openai.com/auth` claim, or nil. The
    /// id_token is a JWT — decode the (unverified) payload segment; Codex trusts the same claim.
    static func accountID(fromIDToken idToken: String) -> String? {
        let segments = idToken.split(separator: ".")
        guard segments.count >= 2,
              let payload = base64URLDecode(String(segments[1])),
              let claims = try? JSONSerialization.jsonObject(with: payload) as? [String: Any]
        else { return nil }
        if let auth = claims["https://api.openai.com/auth"] as? [String: Any],
           let id = auth["chatgpt_account_id"] as? String, !id.isEmpty {
            return id
        }
        // Some tokens carry the id at the top level.
        if let id = claims["chatgpt_account_id"] as? String, !id.isEmpty { return id }
        return nil
    }

    // MARK: - Helpers

    /// ISO-8601 UTC with fractional seconds, matching the `last_refresh` format Codex writes.
    static func iso8601(_ date: Date) -> String {
        let formatter = ISO8601DateFormatter()
        formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
        return formatter.string(from: date)
    }

    static func formURLEncode(_ form: [String: String]) -> String {
        var allowed = CharacterSet.alphanumerics
        allowed.insert(charactersIn: "-._~")
        return form
            .sorted { $0.key < $1.key }
            .map { key, value in
                let k = key.addingPercentEncoding(withAllowedCharacters: allowed) ?? key
                let v = value.addingPercentEncoding(withAllowedCharacters: allowed) ?? value
                return "\(k)=\(v)"
            }
            .joined(separator: "&")
    }

    /// URL-safe base64 without padding (RFC 7636 §A).
    static func base64URL(_ data: Data) -> String {
        data.base64EncodedString()
            .replacingOccurrences(of: "+", with: "-")
            .replacingOccurrences(of: "/", with: "_")
            .replacingOccurrences(of: "=", with: "")
    }

    /// Decode a base64url string (JWT segments have no padding and use the URL-safe alphabet).
    static func base64URLDecode(_ string: String) -> Data? {
        var s = string
            .replacingOccurrences(of: "-", with: "+")
            .replacingOccurrences(of: "_", with: "/")
        while s.count % 4 != 0 { s.append("=") }
        return Data(base64Encoded: s)
    }

    /// The default `Transport`: a `URLSession` round-trip wrapped in a continuation so it compiles
    /// identically on macOS and on the Linux runner.
    public static let liveTransport: Transport = { request in
        try await withCheckedThrowingContinuation { continuation in
            let task = URLSession.shared.dataTask(with: request) { data, response, error in
                if let error {
                    continuation.resume(throwing: error)
                    return
                }
                guard let http = response as? HTTPURLResponse else {
                    continuation.resume(throwing: OAuthError.malformedResponse)
                    return
                }
                continuation.resume(returning: (data ?? Data(), http))
            }
            task.resume()
        }
    }
}
