//
//  OAuthClient.swift
//  AT-IRC
//
//  AT Protocol OAuth for a native public client: PKCE + DPoP + PAR with a
//  loopback redirect (client_id = http://localhost/).
//

import Foundation
import CryptoKit

struct AuthorizationServerMetadata: Decodable {
    let issuer: String
    let authorizationEndpoint: String
    let tokenEndpoint: String
    let pushedAuthorizationRequestEndpoint: String

    enum CodingKeys: String, CodingKey {
        case issuer
        case authorizationEndpoint = "authorization_endpoint"
        case tokenEndpoint = "token_endpoint"
        case pushedAuthorizationRequestEndpoint = "pushed_authorization_request_endpoint"
    }
}

private struct ProtectedResourceMetadata: Decodable {
    let authorizationServers: [String]

    enum CodingKeys: String, CodingKey {
        case authorizationServers = "authorization_servers"
    }
}

private struct PARDocument: Decodable {
    let requestURI: String
    enum CodingKeys: String, CodingKey {
        case requestURI = "request_uri"
    }
}

private struct TokenResponse: Decodable {
    let accessToken: String
    let tokenType: String
    let refreshToken: String?
    let expiresIn: Int?
    let scope: String?
    let sub: String?

    enum CodingKeys: String, CodingKey {
        case accessToken = "access_token"
        case tokenType = "token_type"
        case refreshToken = "refresh_token"
        case expiresIn = "expires_in"
        case scope
        case sub
    }
}

struct OAuthError: LocalizedError {
    let message: String
    var errorDescription: String? { message }
}

@MainActor
final class OAuthClient {
    /// Scopes requested at authorization.
    ///
    /// Exactly the embedded requester's `oauth.RequesterOAuthScope` at the
    /// pinned socialweb-computer-kcp commit (GoLib/build.sh SOCIALWEB_COMMIT;
    /// 13 tokens at 5e6c4e3d, optional scopes included): the market sign-in
    /// reuses this client_id, and the authorization server refuses
    /// (invalid_scope) any requested scope the client_id does not declare.
    /// Note `#` is encoded as `%23` in the scope string itself, and the
    /// client_id query value encodes `%` again.
    static let scope = "atproto repo?action=create&action=update&collection=com.publicdomainrelay.temp.compute.vm&collection=com.publicdomainrelay.temp.market.rfp&collection=com.publicdomainrelay.temp.market.accept&collection=com.publicdomainrelay.temp.market.event&collection=com.publicdomainrelay.temp.compute.events.vm.delete&collection=com.publicdomainrelay.temp.compute.events.vm.onNetwork&collection=com.publicdomainrelay.temp.compute.events.guest.publicKey&collection=com.publicdomainrelay.temp.compute.events.guest.signalWebRTC&collection=com.publicdomainrelay.temp.compute.events.guest.secrets&collection=com.publicdomainrelay.temp.compute.events.guest.agentStart&collection=com.publicdomainrelay.temp.compute.events.guest.secretsUpdated&collection=com.publicdomainrelay.temp.compute.events.guest.r2Start&collection=com.publicdomainrelay.temp.compute.events.guest.r2Release&collection=computer.socialweb.temp.policy.ghalite&collection=computer.socialweb.temp.policy.typescript repo?action=create&action=update&action=delete&collection=com.fedproxy.rbac&collection=com.publicdomainrelay.temp.badgeBlueKeys&collection=com.publicdomainrelay.temp.market.association&collection=com.publicdomainrelay.temp.market.associationRequest&collection=computer.socialweb.infra&collection=computer.socialweb.ephemeralInfra rpc:com.publicdomainrelay.temp.market.submitRfp?aud=* rpc:com.publicdomainrelay.temp.market.submitAccept?aud=* rpc:com.publicdomainrelay.temp.market.submitBid?aud=* rpc:com.publicdomainrelay.temp.market.submitEvent?aud=* rpc:chat.atirc.mint?aud=did:web:irc.fedfork.com%23atirc rpc:chat.atirc.revoke?aud=did:web:irc.fedfork.com%23atirc repo:sh.tangled.graph.vouch?action=create&action=update repo?action=delete&collection=com.publicdomainrelay.temp.compute.vm&collection=com.publicdomainrelay.temp.market.rfp&collection=com.publicdomainrelay.temp.market.accept&collection=com.publicdomainrelay.temp.market.event&collection=com.publicdomainrelay.temp.compute.events.vm.delete&collection=com.publicdomainrelay.temp.compute.events.vm.onNetwork&collection=com.publicdomainrelay.temp.compute.events.guest.publicKey&collection=com.publicdomainrelay.temp.compute.events.guest.signalWebRTC&collection=com.publicdomainrelay.temp.compute.events.guest.secrets&collection=com.publicdomainrelay.temp.compute.events.guest.agentStart&collection=com.publicdomainrelay.temp.compute.events.guest.secretsUpdated&collection=com.publicdomainrelay.temp.compute.events.guest.r2Start&collection=com.publicdomainrelay.temp.compute.events.guest.r2Release&collection=computer.socialweb.temp.policy.ghalite&collection=computer.socialweb.temp.policy.typescript repo:computer.socialweb.ingress.offer?action=create&action=delete repo:computer.socialweb.peer?action=create&action=update&action=delete"

    /// For loopback clients the authorization server synthesizes the client
    /// metadata; any scope beyond `atproto` must be declared in the client_id
    /// URL's `scope` query parameter.
    static var clientID: String {
        var components = URLComponents(string: "http://localhost/")!
        components.queryItems = [URLQueryItem(name: "scope", value: scope)]
        return components.url?.absoluteString ?? "http://localhost/"
    }

    private var nonces: [String: String] = [:]

    // MARK: - Public flow

    struct PendingAuthorization {
        let state: String
        let verifier: String
        let redirectURI: String
        let dpopKey: DPoPKey
        let metadata: AuthorizationServerMetadata
        let pdsURL: URL
        let did: String
        let handle: String?
        let loopback: LoopbackServer
    }

    /// Resolve the account, start the loopback listener, and push the
    /// authorization request. Returns the URL to present in a browser.
    func beginSignIn(handle input: String) async throws -> (url: URL, pending: PendingAuthorization) {
        let identity = try await resolveIdentity(input)
        let metadata = try await discoverAuthorizationServer(pds: identity.pds)

        let loopback = LoopbackServer()
        let port = try await loopback.start()
        let redirectURI = "http://127.0.0.1:\(port)/"

        let verifier = Base64URL.encode(.random(count: 32))
        let challenge = Base64URL.encode(Data(SHA256.hash(data: Data(verifier.utf8))))
        let state = Base64URL.encode(.random(count: 16))
        let key = DPoPKey()

        var form: [String: String] = [
            "client_id": Self.clientID,
            "response_type": "code",
            "code_challenge": challenge,
            "code_challenge_method": "S256",
            "redirect_uri": redirectURI,
            "scope": Self.scope,
            "state": state,
        ]
        if !identity.handle.isEmpty {
            form["login_hint"] = identity.handle
        }

        let parURL = try url(metadata.pushedAuthorizationRequestEndpoint)
        let (data, response) = try await sendDPoP(
            url: parURL,
            method: "POST",
            form: form,
            key: key,
            accessToken: nil,
            nonceKey: parURL.host ?? "par"
        )
        guard (200..<300).contains(response.statusCode) else {
            throw OAuthError(message: Self.describe(status: response.statusCode, data: data))
        }
        let par = try JSONDecoder().decode(PARDocument.self, from: data)

        var components = URLComponents(string: metadata.authorizationEndpoint)
        components?.queryItems = [
            URLQueryItem(name: "client_id", value: Self.clientID),
            URLQueryItem(name: "request_uri", value: par.requestURI),
        ]
        guard let authorizeURL = components?.url else {
            throw OAuthError(message: "Bad authorization endpoint")
        }

        let pending = PendingAuthorization(
            state: state,
            verifier: verifier,
            redirectURI: redirectURI,
            dpopKey: key,
            metadata: metadata,
            pdsURL: identity.pds,
            did: identity.did,
            handle: identity.handle.isEmpty ? nil : identity.handle,
            loopback: loopback
        )
        return (authorizeURL, pending)
    }

    /// Exchange the authorization code after the browser redirects back.
    func completeSignIn(pending: PendingAuthorization, query: [String: String]) async throws -> OAuthSession {
        defer { pending.loopback.stop() }

        if let error = query["error"] {
            let description = query["error_description"] ?? error
            throw OAuthError(message: "Authorization failed: \(description)")
        }
        guard query["state"] == pending.state else {
            throw OAuthError(message: "OAuth state mismatch")
        }
        if let issuer = query["iss"], issuer != pending.metadata.issuer {
            throw OAuthError(message: "Authorization server mismatch")
        }
        guard let code = query["code"] else {
            throw OAuthError(message: "No authorization code returned")
        }

        let tokenURL = try url(pending.metadata.tokenEndpoint)
        let form: [String: String] = [
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": pending.redirectURI,
            "client_id": Self.clientID,
            "code_verifier": pending.verifier,
        ]
        let (data, response) = try await sendDPoP(
            url: tokenURL,
            method: "POST",
            form: form,
            key: pending.dpopKey,
            accessToken: nil,
            nonceKey: tokenURL.host ?? "token"
        )
        guard (200..<300).contains(response.statusCode) else {
            throw OAuthError(message: Self.describe(status: response.statusCode, data: data))
        }
        let token = try JSONDecoder().decode(TokenResponse.self, from: data)
        if let scope = token.scope, !scope.split(separator: " ").contains("atproto") {
            throw OAuthError(message: "Token is not an atproto token")
        }
        if let sub = token.sub, sub != pending.did {
            throw OAuthError(message: "Signed in as an unexpected account")
        }

        return OAuthSession(
            accessToken: token.accessToken,
            refreshToken: token.refreshToken,
            expiresAt: Date().addingTimeInterval(TimeInterval(token.expiresIn ?? 300)),
            tokenType: token.tokenType,
            did: token.sub ?? pending.did,
            handle: pending.handle,
            pdsURL: pending.pdsURL.absoluteString,
            authServer: pending.metadata.issuer,
            dpopKey: pending.dpopKey.rawRepresentation.base64EncodedString(),
            scope: token.scope
        )
    }

    /// Refresh an expired access token (refresh tokens are single-use).
    func refresh(_ session: OAuthSession) async throws -> OAuthSession {
        guard let refreshToken = session.refreshToken, let key = session.dpopKeyValue else {
            throw OAuthError(message: "Session cannot be refreshed")
        }
        let metadata = try await discoverAuthorizationServer(pds: try url(session.pdsURL))
        let tokenURL = try url(metadata.tokenEndpoint)
        let form: [String: String] = [
            "grant_type": "refresh_token",
            "refresh_token": refreshToken,
            "client_id": Self.clientID,
        ]
        let (data, response) = try await sendDPoP(
            url: tokenURL,
            method: "POST",
            form: form,
            key: key,
            accessToken: nil,
            nonceKey: tokenURL.host ?? "token"
        )
        guard (200..<300).contains(response.statusCode) else {
            throw OAuthError(message: Self.describe(status: response.statusCode, data: data))
        }
        let token = try JSONDecoder().decode(TokenResponse.self, from: data)

        var updated = session
        updated.accessToken = token.accessToken
        updated.refreshToken = token.refreshToken ?? session.refreshToken
        updated.expiresAt = Date().addingTimeInterval(TimeInterval(token.expiresIn ?? 300))
        updated.scope = token.scope ?? session.scope
        updated.authServer = metadata.issuer
        return updated
    }

    /// Ask the PDS for a service-auth JWT bound to `lxm` and the atirc service DID.
    func serviceAuth(session: OAuthSession, lxm: String, aud: String) async throws -> (token: String, session: OAuthSession) {
        var current = session
        if current.isExpired {
            current = try await refresh(current)
        }
        guard let key = current.dpopKeyValue else {
            throw OAuthError(message: "Missing DPoP key")
        }
        var components = URLComponents(string: current.pdsURL.trimmingCharacters(in: CharacterSet(charactersIn: "/")) + "/xrpc/com.atproto.server.getServiceAuth")!
        let expiry = Int(Date().timeIntervalSince1970) + 120
        components.queryItems = [
            URLQueryItem(name: "aud", value: aud),
            URLQueryItem(name: "lxm", value: lxm),
            URLQueryItem(name: "exp", value: String(expiry)),
        ]
        guard let requestURL = components.url else {
            throw OAuthError(message: "Bad getServiceAuth URL")
        }
        let (data, response) = try await sendDPoP(
            url: requestURL,
            method: "GET",
            form: nil,
            key: key,
            accessToken: current.accessToken,
            nonceKey: requestURL.host ?? "pds"
        )
        if response.statusCode == 401, current.refreshToken != nil {
            // Access token was rejected; refresh once and retry.
            let refreshed = try await refresh(current)
            return try await serviceAuth(session: refreshed, lxm: lxm, aud: aud)
        }
        guard (200..<300).contains(response.statusCode) else {
            throw OAuthError(message: Self.describe(status: response.statusCode, data: data))
        }
        struct ServiceAuthResponse: Decodable { let token: String }
        let decoded = try JSONDecoder().decode(ServiceAuthResponse.self, from: data)
        return (decoded.token, current)
    }

    // MARK: - Identity + discovery

    struct ResolvedIdentity {
        let did: String
        let handle: String
        let pds: URL
    }

    func resolveIdentity(_ input: String) async throws -> ResolvedIdentity {
        var value = input.trimmingCharacters(in: .whitespacesAndNewlines)
        if value.hasPrefix("@") { value.removeFirst() }
        guard !value.isEmpty else { throw OAuthError(message: "Enter a handle") }

        var handle: String?
        let did: String
        if value.lowercased().hasPrefix("did:") {
            did = value
        } else {
            guard let resolveURL = URL(string: "https://public.api.bsky.app/xrpc/com.atproto.identity.resolveHandle?handle=\(value.addingPercentEncoding(withAllowedCharacters: .urlQueryAllowed) ?? value)") else {
                throw OAuthError(message: "Invalid handle")
            }
            let (data, response) = try await URLSession.shared.data(from: resolveURL)
            guard let http = response as? HTTPURLResponse, (200..<300).contains(http.statusCode) else {
                throw OAuthError(message: "Could not resolve handle \(value)")
            }
            struct ResolveResponse: Decodable { let did: String }
            did = try JSONDecoder().decode(ResolveResponse.self, from: data).did
            handle = value
        }

        let document = try await fetchDIDDocument(did)
        guard let pds = Self.pdsEndpoint(from: document) else {
            throw OAuthError(message: "No PDS found in DID document")
        }
        return ResolvedIdentity(did: did, handle: handle ?? Self.handle(from: document) ?? "", pds: pds)
    }

    func discoverAuthorizationServer(pds: URL) async throws -> AuthorizationServerMetadata {
        let base = pds.absoluteString.trimmingCharacters(in: CharacterSet(charactersIn: "/"))
        let resourceURL = try url(base + "/.well-known/oauth-protected-resource")
        var authServer: URL
        do {
            let (data, response) = try await URLSession.shared.data(from: resourceURL)
            guard let http = response as? HTTPURLResponse, (200..<300).contains(http.statusCode) else {
                throw OAuthError(message: "No protected resource metadata")
            }
            let resource = try JSONDecoder().decode(ProtectedResourceMetadata.self, from: data)
            guard let first = resource.authorizationServers.first, let url = URL(string: first) else {
                throw OAuthError(message: "No authorization server advertised")
            }
            authServer = url
        } catch {
            // Fall back to the PDS itself as the authorization server.
            authServer = pds
        }

        let metadataURL = try url(authServer.absoluteString.trimmingCharacters(in: CharacterSet(charactersIn: "/")) + "/.well-known/oauth-authorization-server")
        let (data, response) = try await URLSession.shared.data(from: metadataURL)
        guard let http = response as? HTTPURLResponse, (200..<300).contains(http.statusCode) else {
            throw OAuthError(message: "Could not load authorization server metadata")
        }
        let metadata = try JSONDecoder().decode(AuthorizationServerMetadata.self, from: data)
        guard metadata.issuer == authServer.absoluteString.trimmingCharacters(in: CharacterSet(charactersIn: "/")) ||
              URL(string: metadata.issuer)?.host == authServer.host else {
            throw OAuthError(message: "Authorization server issuer mismatch")
        }
        return metadata
    }

    private func fetchDIDDocument(_ did: String) async throws -> [String: Any] {
        let urlString: String
        if did.hasPrefix("did:plc:") {
            urlString = "https://plc.directory/\(did)"
        } else if did.hasPrefix("did:web:") {
            let host = did.dropFirst("did:web:".count).replacingOccurrences(of: "%3A", with: ":")
            urlString = "https://\(host)/.well-known/did.json"
        } else {
            throw OAuthError(message: "Unsupported DID method")
        }
        let (data, response) = try await URLSession.shared.data(from: try url(urlString))
        guard let http = response as? HTTPURLResponse, (200..<300).contains(http.statusCode) else {
            throw OAuthError(message: "Could not resolve DID document")
        }
        guard let object = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
            throw OAuthError(message: "Malformed DID document")
        }
        return object
    }

    private static func pdsEndpoint(from document: [String: Any]) -> URL? {
        guard let services = document["service"] as? [[String: Any]] else { return nil }
        for service in services {
            let id = service["id"] as? String ?? ""
            let type = service["type"] as? String ?? ""
            if id == "#atproto_pds" || type == "AtprotoPersonalDataServer" {
                if let endpoint = service["serviceEndpoint"] as? String {
                    return URL(string: endpoint)
                }
            }
        }
        return nil
    }

    private static func handle(from document: [String: Any]) -> String? {
        document["alsoKnownAs"] as? [String] == nil ? nil : (document["alsoKnownAs"] as? [String])?
            .first(where: { $0.hasPrefix("at://") })?
            .replacingOccurrences(of: "at://", with: "")
    }

    // MARK: - HTTP

    private func sendDPoP(
        url: URL,
        method: String,
        form: [String: String]?,
        key: DPoPKey,
        accessToken: String?,
        nonceKey: String
    ) async throws -> (Data, HTTPURLResponse) {
        var request = URLRequest(url: url)
        request.httpMethod = method
        request.timeoutInterval = 30
        if let form {
            request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type")
            request.httpBody = Data(Self.formEncode(form).utf8)
        }
        request.setValue("application/json", forHTTPHeaderField: "Accept")

        var lastData = Data()
        var lastResponse: HTTPURLResponse?
        for attempt in 0...1 {
            let proof = try key.proof(
                method: method,
                url: url,
                nonce: nonces[nonceKey],
                accessToken: accessToken
            )
            request.setValue(proof, forHTTPHeaderField: "DPoP")
            if let accessToken {
                request.setValue("DPoP \(accessToken)", forHTTPHeaderField: "Authorization")
            }
            let (data, response) = try await URLSession.shared.data(for: request)
            guard let http = response as? HTTPURLResponse else {
                throw OAuthError(message: "Bad server response")
            }
            if let nonce = http.value(forHTTPHeaderField: "DPoP-Nonce") {
                nonces[nonceKey] = nonce
            }
            lastData = data
            lastResponse = http

            let needsNonce = http.statusCode == 400 &&
                (String(data: data, encoding: .utf8) ?? "").contains("use_dpop_nonce")
            if needsNonce, attempt == 0, nonces[nonceKey] != nil {
                continue
            }
            return (data, http)
        }
        return (lastData, lastResponse ?? HTTPURLResponse())
    }

    private func url(_ string: String) throws -> URL {
        guard let url = URL(string: string) else {
            throw OAuthError(message: "Invalid URL: \(string)")
        }
        return url
    }

    private static func formEncode(_ values: [String: String]) -> String {
        values
            .sorted { $0.key < $1.key }
            .map { key, value in
                let escapedKey = key.addingPercentEncoding(withAllowedCharacters: .alphanumerics) ?? key
                let escapedValue = value.addingPercentEncoding(withAllowedCharacters: .alphanumerics) ?? value
                return "\(escapedKey)=\(escapedValue)"
            }
            .joined(separator: "&")
    }

    private static func describe(status: Int, data: Data) -> String {
        if let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
           let message = object["error_description"] as? String ?? object["message"] as? String ?? object["error"] as? String {
            return "HTTP \(status): \(message)"
        }
        let text = String(data: data, encoding: .utf8) ?? ""
        return "HTTP \(status)\(text.isEmpty ? "" : ": \(text.prefix(200))")"
    }
}