import Foundation
import CryptoKit

enum AiConfigClient {
    struct PairPayload {
        let baseURL: String
        let pairId: String
        let serverPubB64: String
    }

    struct Meta: Decodable {
        let version: UInt64
        let updatedAt: String
        let source: String
    }

    struct EncryptedConfig: Decodable {
        let version: UInt64
        let updatedAt: String
        let source: String
        let nonce: String
        let ciphertext: String
    }

    static func parsePairPayload(_ raw: String) -> PairPayload? {
        let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines)
        guard trimmed.hasPrefix("studydeck-ai:1?") else {
            // Also accept plain base URL for health-only (not pair)
            return nil
        }
        let q = String(trimmed.dropFirst("studydeck-ai:1?".count))
        var map: [String: String] = [:]
        for part in q.split(separator: "&") {
            let kv = part.split(separator: "=", maxSplits: 1).map(String.init)
            guard kv.count == 2 else { continue }
            map[kv[0]] = kv[1].removingPercentEncoding ?? kv[1]
        }
        guard let base = map["base"], let pairId = map["pairId"], let pub = map["pub"] else {
            return nil
        }
        return PairPayload(baseURL: base, pairId: pairId, serverPubB64: pub)
    }

    static func pair(payload: PairPayload, deviceName: String = "iPad") async throws {
        let clientPrivate = Curve25519.KeyAgreement.PrivateKey()
        let clientPubB64 = Data(clientPrivate.publicKey.rawRepresentation).base64EncodedString()

        guard let serverPubData = Data(base64Encoded: payload.serverPubB64),
              let serverPub = try? Curve25519.KeyAgreement.PublicKey(rawRepresentation: serverPubData)
        else {
            throw URLError(.badServerResponse)
        }

        let shared = try clientPrivate.sharedSecretFromKeyAgreement(with: serverPub)
        let aesKey = shared.hkdfDerivedSymmetricKey(
            using: SHA256.self,
            salt: Data(payload.pairId.utf8),
            sharedInfo: Data("studydeck-ai-config-v1".utf8),
            outputByteCount: 32
        )

        let url = URL(string: payload.baseURL.trimmingCharacters(in: CharacterSet(charactersIn: "/")) + "/ai/pair")!
        var req = URLRequest(url: url)
        req.httpMethod = "POST"
        req.setValue("application/json", forHTTPHeaderField: "Content-Type")
        req.timeoutInterval = 15
        let body: [String: String] = [
            "pairId": payload.pairId,
            "clientPub": clientPubB64,
            "deviceName": deviceName,
        ]
        req.httpBody = try JSONSerialization.data(withJSONObject: body)
        let (data, resp) = try await URLSession.shared.data(for: req)
        guard let http = resp as? HTTPURLResponse, (200..<300).contains(http.statusCode) else {
            throw URLError(.badServerResponse)
        }
        let json = try JSONSerialization.jsonObject(with: data) as? [String: Any]
        guard let token = json?["sessionToken"] as? String else {
            throw URLError(.cannotParseResponse)
        }
        AiConfigStore.saveSession(
            baseURL: payload.baseURL,
            token: token,
            aesKey: aesKey.withUnsafeBytes { Data($0) }
        )
        _ = try await syncFromPC(force: true)
    }

    static func fetchMeta() async throws -> Meta {
        guard let base = AiConfigStore.shareBaseURL, let token = AiConfigStore.sessionToken else {
            throw AiConfigResolveError.unpaired
        }
        var req = URLRequest(url: URL(string: base.trimmingCharacters(in: CharacterSet(charactersIn: "/")) + "/ai/config/meta")!)
        req.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
        req.timeoutInterval = 8
        let (data, resp) = try await URLSession.shared.data(for: req)
        if let http = resp as? HTTPURLResponse, http.statusCode == 401 {
            throw AiConfigResolveError.stale
        }
        guard let http = resp as? HTTPURLResponse, (200..<300).contains(http.statusCode) else {
            throw URLError(.badServerResponse)
        }
        return try JSONDecoder().decode(Meta.self, from: data)
    }

    static func fetchEncryptedConfig() async throws -> AiRuntimeConfig {
        guard let base = AiConfigStore.shareBaseURL,
              let token = AiConfigStore.sessionToken,
              let aesRaw = AiConfigStore.aesKey
        else {
            throw AiConfigResolveError.unpaired
        }
        var req = URLRequest(url: URL(string: base.trimmingCharacters(in: CharacterSet(charactersIn: "/")) + "/ai/config")!)
        req.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
        req.timeoutInterval = 20
        let (data, resp) = try await URLSession.shared.data(for: req)
        if let http = resp as? HTTPURLResponse, http.statusCode == 401 {
            throw AiConfigResolveError.stale
        }
        guard let http = resp as? HTTPURLResponse, (200..<300).contains(http.statusCode) else {
            throw URLError(.badServerResponse)
        }
        let enc = try JSONDecoder().decode(EncryptedConfig.self, from: data)
        guard let nonceData = Data(base64Encoded: enc.nonce),
              let ct = Data(base64Encoded: enc.ciphertext)
        else {
            throw URLError(.cannotParseResponse)
        }
        // Rust aes-gcm returns ciphertext||tag; CryptoKit combined = nonce||ciphertext||tag
        let key = SymmetricKey(data: aesRaw)
        let sealed = try AES.GCM.SealedBox(combined: nonceData + ct)
        let plain = try AES.GCM.open(sealed, using: key)
        let cfg = try JSONDecoder().decode(AiRuntimeConfig.self, from: plain)
        try AiConfigStore.saveConfig(cfg)
        return cfg
    }

    /// Resolve config: PC online → compare updatedAt; offline → cache.
    static func resolveConfig(preferFeature: ((AiRuntimeConfig) -> Bool)? = nil) async throws -> AiRuntimeConfig {
        var candidates: [AiRuntimeConfig] = []
        if let cached = AiConfigStore.loadConfig() {
            candidates.append(cached)
        }

        if AiConfigStore.isPaired {
            do {
                let meta = try await fetchMeta()
                let localAt = AiConfigStore.cachedUpdatedAt ?? ""
                if meta.updatedAt > localAt || AiConfigStore.loadConfig() == nil {
                    let fresh = try await fetchEncryptedConfig()
                    candidates = [fresh]
                }
            } catch AiConfigResolveError.stale {
                AiConfigStore.clearSession()
                if candidates.isEmpty { throw AiConfigResolveError.stale }
            } catch {
                // PC unreachable — fall through to cache
                if candidates.isEmpty {
                    throw AiConfigResolveError.shareUnavailable
                }
            }
        } else if candidates.isEmpty {
            throw AiConfigResolveError.unpaired
        }

        // Phase C stub: server source merge (disabled unless serverSyncEnabled on cached config)
        if let local = candidates.first, local.serverSyncEnabled, !local.serverBaseUrl.isEmpty {
            if let server = await fetchServerConfigStub(local) {
                candidates.append(server)
            }
        }

        guard let best = candidates.max(by: { $0.updatedAt < $1.updatedAt }) else {
            throw AiConfigResolveError.noCache
        }
        let hasKey = !best.apiKey.isEmpty || best.endpoints.contains(where: { !$0.apiKey.isEmpty })
        if !hasKey {
            throw AiConfigResolveError.noCache
        }
        if let preferFeature, !preferFeature(best) {
            throw AiConfigResolveError.featureDisabled
        }
        return best
    }

    static func syncFromPC(force: Bool) async throws -> AiRuntimeConfig {
        if force {
            return try await fetchEncryptedConfig()
        }
        return try await resolveConfig()
    }

    /// Reserved: subscription server. Returns nil when not configured / unimplemented.
    static func fetchServerConfigStub(_ local: AiRuntimeConfig) async -> AiRuntimeConfig? {
        _ = local
        return nil
    }

    static func probeHealth(baseURL: String) async -> Bool {
        guard let url = URL(string: baseURL.trimmingCharacters(in: CharacterSet(charactersIn: "/")) + "/ai/health") else {
            return false
        }
        var req = URLRequest(url: url)
        req.timeoutInterval = 5
        do {
            let (_, resp) = try await URLSession.shared.data(for: req)
            return (resp as? HTTPURLResponse)?.statusCode == 200
        } catch {
            return false
        }
    }
}