import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif

#if os(macOS)
import SweetCookieKit
#endif

public enum AmpUsageError: LocalizedError, Sendable {
    case notLoggedIn
    case invalidCredentials
    case missingAPIToken
    case invalidAPIToken
    case parseFailed(String)
    case networkError(String)
    case noSessionCookie

    public var errorDescription: String? {
        switch self {
        case .notLoggedIn:
            "Not logged in to Amp. Please log in via ampcode.com."
        case .invalidCredentials:
            "Amp session cookie expired. Please log in again."
        case .missingAPIToken:
            "Amp access token not configured. Set AMP_API_KEY or add it in Settings."
        case .invalidAPIToken:
            "Amp access token is invalid or expired."
        case let .parseFailed(message):
            "Could not parse Amp usage: \(message)"
        case let .networkError(message):
            "Amp request failed: \(message)"
        case .noSessionCookie:
            "No Amp session cookie found. Please log in to ampcode.com in your browser."
        }
    }
}

#if os(macOS)
private let ampCookieImportOrder: BrowserCookieImportOrder =
    ProviderDefaults.metadata[.amp]?.browserCookieOrder ?? Browser.defaultImportOrder

public enum AmpCookieImporter {
    private static let cookieClient = BrowserCookieClient()
    private static let cookieDomains = ["ampcode.com", "www.ampcode.com"]
    private static let sessionCookieNames: Set<String> = [
        "session",
    ]

    public struct SessionInfo: Sendable {
        public let cookies: [HTTPCookie]
        public let sourceLabel: String

        public init(cookies: [HTTPCookie], sourceLabel: String) {
            self.cookies = cookies
            self.sourceLabel = sourceLabel
        }

        public var cookieHeader: String {
            self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ")
        }
    }

    public static func importSession(
        browserDetection: BrowserDetection,
        logger: ((String) -> Void)? = nil) throws -> SessionInfo
    {
        let log: (String) -> Void = { msg in logger?("[amp-cookie] \(msg)") }

        let installed = ampCookieImportOrder.cookieImportCandidates(using: browserDetection)
        for browserSource in installed {
            do {
                let query = BrowserCookieQuery(domains: self.cookieDomains)
                let sources = try Self.cookieClient.codexBarRecords(
                    matching: query,
                    in: browserSource,
                    logger: log)
                for source in sources where !source.records.isEmpty {
                    let cookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin)
                    guard !cookies.isEmpty else { continue }
                    let names = cookies.map(\.name).joined(separator: ", ")
                    log("\(source.label) cookies: \(names)")
                    let sessionCookies = cookies.filter { Self.sessionCookieNames.contains($0.name) }
                    if !sessionCookies.isEmpty {
                        log("Found Amp session cookie in \(source.label)")
                        return SessionInfo(cookies: sessionCookies, sourceLabel: source.label)
                    }
                    log("\(source.label) cookies found, but no Amp session cookie present")
                    log("Expected one of: \(Self.sessionCookieNames.joined(separator: ", "))")
                }
            } catch {
                BrowserCookieAccessGate.recordIfNeeded(error)
                log("\(browserSource.displayName) cookie import failed: \(error.localizedDescription)")
            }
        }

        throw AmpUsageError.noSessionCookie
    }
}
#endif

public struct AmpUsageFetcher: Sendable {
    private static let settingsURL = URL(string: "https://ampcode.com/settings")!
    static let usageURL = URL(string: "https://ampcode.com/api/internal?userDisplayBalanceInfo")!
    @MainActor private static var recentDumps: [String] = []

    public let browserDetection: BrowserDetection

    public init(browserDetection: BrowserDetection) {
        self.browserDetection = browserDetection
    }

    public func fetch(
        cookieHeaderOverride: String? = nil,
        logger: ((String) -> Void)? = nil,
        now: Date = Date()) async throws -> AmpUsageSnapshot
    {
        let log: (String) -> Void = { msg in logger?("[amp] \(msg)") }
        let cookieHeader = try await self.resolveCookieHeader(override: cookieHeaderOverride, logger: log)

        if let logger {
            let names = self.cookieNames(from: cookieHeader)
            if !names.isEmpty {
                logger("[amp] Cookie names: \(names.joined(separator: ", "))")
            }
            let diagnostics = RedirectDiagnostics(cookieHeader: cookieHeader, logger: logger)
            do {
                let (html, responseInfo) = try await self.fetchLegacyHTMLWithDiagnostics(
                    cookieHeader: cookieHeader,
                    diagnostics: diagnostics)
                self.logDiagnostics(responseInfo: responseInfo, diagnostics: diagnostics, logger: logger)
                return try AmpUsageParser.parse(html: html, now: now)
            } catch {
                self.logDiagnostics(responseInfo: nil, diagnostics: diagnostics, logger: logger)
                logger("[amp] Fetch failed: \(error.localizedDescription)")
                throw error
            }
        }

        let diagnostics = RedirectDiagnostics(cookieHeader: cookieHeader, logger: nil)
        let (html, _) = try await self.fetchLegacyHTMLWithDiagnostics(
            cookieHeader: cookieHeader,
            diagnostics: diagnostics)
        return try AmpUsageParser.parse(html: html, now: now)
    }

    public func fetch(
        apiToken: String,
        logger: ((String) -> Void)? = nil,
        now: Date = Date()) async throws -> AmpUsageSnapshot
    {
        guard let token = AmpSettingsReader.cleaned(apiToken) else {
            throw AmpUsageError.missingAPIToken
        }
        let request = try Self.makeUsageAPIRequest(apiToken: token)
        let diagnostics = APIRedirectDiagnostics(logger: logger)
        let session = URLSession(configuration: .ephemeral, delegate: diagnostics, delegateQueue: nil)
        let httpResponse = try await session.response(for: request)
        logger?("[amp] API response: \(httpResponse.statusCode) " +
            "\(httpResponse.response.url?.absoluteString ?? "unknown")")
        try Self.validateAPIResponse(httpResponse)
        return try Self.parseUsageAPIResponse(httpResponse.data, now: now)
    }

    public func debugRawProbe(cookieHeaderOverride: String? = nil) async -> String {
        let stamp = ISO8601DateFormatter().string(from: Date())
        var lines: [String] = []
        lines.append("=== Amp Debug Probe @ \(stamp) ===")
        lines.append("")

        do {
            let cookieHeader = try await self.resolveCookieHeader(
                override: cookieHeaderOverride,
                logger: { msg in lines.append("[cookie] \(msg)") })
            let diagnostics = RedirectDiagnostics(cookieHeader: cookieHeader, logger: nil)
            let cookieNames = CookieHeaderNormalizer.pairs(from: cookieHeader).map(\.name)
            lines.append("Cookie names: \(cookieNames.joined(separator: ", "))")

            let (html, responseInfo) = try await self.fetchLegacyHTMLWithDiagnostics(
                cookieHeader: cookieHeader,
                diagnostics: diagnostics)
            let snapshot = try AmpUsageParser.parse(html: html)

            lines.append("")
            lines.append("Fetch Success")
            lines.append("Status: \(responseInfo.statusCode) \(responseInfo.url)")

            if !diagnostics.redirects.isEmpty {
                lines.append("")
                lines.append("Redirects:")
                for entry in diagnostics.redirects {
                    lines.append("  \(entry)")
                }
            }

            lines.append("")
            lines.append("Amp Free:")
            lines.append("  quota=\(snapshot.freeQuota?.description ?? "nil")")
            lines.append("  used=\(snapshot.freeUsed?.description ?? "nil")")
            lines.append("  hourlyReplenishment=\(snapshot.hourlyReplenishment?.description ?? "nil")")
            lines.append("  windowHours=\(snapshot.windowHours?.description ?? "nil")")
            lines.append("  individualCredits=\(snapshot.individualCredits?.description ?? "nil")")
            for workspace in snapshot.workspaceBalances {
                lines.append("  workspace[\(workspace.name)]=\(workspace.remaining)")
            }

            let output = lines.joined(separator: "\n")
            Task { @MainActor in Self.recordDump(output) }
            return output
        } catch {
            lines.append("")
            lines.append("Probe Failed: \(error.localizedDescription)")
            let output = lines.joined(separator: "\n")
            Task { @MainActor in Self.recordDump(output) }
            return output
        }
    }

    public static func latestDumps() async -> String {
        await MainActor.run {
            let result = Self.recentDumps.joined(separator: "\n\n---\n\n")
            return result.isEmpty ? "No Amp probe dumps captured yet." : result
        }
    }

    private func resolveCookieHeader(
        override: String?,
        logger: ((String) -> Void)?) async throws -> String
    {
        if let override = CookieHeaderNormalizer.normalize(override) {
            if let sessionHeader = self.sessionCookieHeader(from: override) {
                logger?("[amp] Using manual session cookie")
                return sessionHeader
            }
            throw AmpUsageError.noSessionCookie
        }
        #if os(macOS)
        let session = try AmpCookieImporter.importSession(browserDetection: self.browserDetection, logger: logger)
        logger?("[amp] Using cookies from \(session.sourceLabel)")
        return session.cookieHeader
        #else
        throw AmpUsageError.noSessionCookie
        #endif
    }

    static func makeUsageAPIRequest(apiToken: String) throws -> URLRequest {
        var request = URLRequest(url: Self.usageURL)
        request.httpMethod = "POST"
        request.httpBody = try JSONSerialization.data(withJSONObject: [
            "method": "userDisplayBalanceInfo",
            "params": [:],
        ])
        request.setValue("Bearer \(apiToken)", forHTTPHeaderField: "Authorization")
        request.setValue("application/json", forHTTPHeaderField: "accept")
        request.setValue("application/json", forHTTPHeaderField: "content-type")
        return request
    }

    private func fetchLegacyHTMLWithDiagnostics(
        cookieHeader: String,
        diagnostics: RedirectDiagnostics) async throws -> (String, ResponseInfo)
    {
        var request = URLRequest(url: Self.settingsURL)
        request.httpMethod = "GET"
        request.setValue(cookieHeader, forHTTPHeaderField: "Cookie")
        request.setValue(
            "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
            forHTTPHeaderField: "accept")
        Self.applyBrowserHeaders(to: &request)

        let session = URLSession(configuration: .ephemeral, delegate: diagnostics, delegateQueue: nil)
        let httpResponse = try await session.response(for: request)
        let responseInfo = ResponseInfo(
            statusCode: httpResponse.statusCode,
            url: httpResponse.response.url?.absoluteString ?? "unknown")
        try Self.validateBrowserResponse(response: httpResponse, diagnostics: diagnostics)

        let html = String(data: httpResponse.data, encoding: .utf8) ?? ""
        return (html, responseInfo)
    }

    static func parseUsageAPIResponse(_ data: Data, now: Date = Date()) throws -> AmpUsageSnapshot {
        let response: UsageAPIResponse
        do {
            response = try JSONDecoder().decode(UsageAPIResponse.self, from: data)
        } catch {
            throw AmpUsageError.parseFailed("Invalid Amp usage API response.")
        }

        guard response.ok else {
            if response.error?.code == "auth-required" {
                throw AmpUsageError.invalidAPIToken
            }
            throw AmpUsageError.networkError(response.error?.message ?? "Amp usage API returned an error.")
        }
        guard let displayText = response.result?.displayText, !displayText.isEmpty else {
            throw AmpUsageError.parseFailed("Missing Amp usage display text.")
        }
        return try AmpUsageParser.parse(displayText: displayText, now: now)
    }

    private static func applyBrowserHeaders(to request: inout URLRequest) {
        request.setValue(
            "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " +
                "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36",
            forHTTPHeaderField: "user-agent")
        request.setValue("en-US,en;q=0.9", forHTTPHeaderField: "accept-language")
        request.setValue("https://ampcode.com", forHTTPHeaderField: "origin")
        request.setValue(self.settingsURL.absoluteString, forHTTPHeaderField: "referer")
    }

    private static func validateBrowserResponse(
        response: ProviderHTTPResponse,
        diagnostics: RedirectDiagnostics) throws
    {
        guard response.statusCode == 200 else {
            if response.statusCode == 401 || response.statusCode == 403 || diagnostics.detectedLoginRedirect {
                throw AmpUsageError.invalidCredentials
            }
            throw AmpUsageError.networkError("HTTP \(response.statusCode)")
        }
    }

    private static func validateAPIResponse(_ response: ProviderHTTPResponse) throws {
        guard response.statusCode == 200 else {
            if response.statusCode == 401 || response.statusCode == 403 {
                throw AmpUsageError.invalidAPIToken
            }
            throw AmpUsageError.networkError("HTTP \(response.statusCode)")
        }
    }

    @MainActor private static func recordDump(_ text: String) {
        if self.recentDumps.count >= 5 { self.recentDumps.removeFirst() }
        self.recentDumps.append(text)
    }

    private final class RedirectDiagnostics: NSObject, URLSessionTaskDelegate, @unchecked Sendable {
        private let cookieHeader: String
        private let logger: ((String) -> Void)?
        var redirects: [String] = []
        private(set) var detectedLoginRedirect = false

        init(cookieHeader: String, logger: ((String) -> Void)?) {
            self.cookieHeader = cookieHeader
            self.logger = logger
        }

        func urlSession(
            _: URLSession,
            task: URLSessionTask,
            willPerformHTTPRedirection response: HTTPURLResponse,
            newRequest request: URLRequest,
            completionHandler: @escaping (URLRequest?) -> Void)
        {
            let from = response.url?.absoluteString ?? "unknown"
            let to = request.url?.absoluteString ?? "unknown"
            self.redirects.append("\(response.statusCode) \(from) -> \(to)")

            if let toURL = request.url, AmpUsageFetcher.isLoginRedirect(toURL) {
                if let logger {
                    logger("[amp] Detected login redirect, aborting (invalid session)")
                }
                self.detectedLoginRedirect = true
                completionHandler(nil)
                return
            }

            var updated = request
            if AmpUsageFetcher.shouldAttachCookie(to: request.url), !self.cookieHeader.isEmpty {
                updated.setValue(self.cookieHeader, forHTTPHeaderField: "Cookie")
            } else {
                updated.setValue(nil, forHTTPHeaderField: "Cookie")
            }
            if let referer = response.url?.absoluteString {
                updated.setValue(referer, forHTTPHeaderField: "referer")
            }
            if let logger {
                logger("[amp] Redirect \(response.statusCode) \(from) -> \(to)")
            }
            completionHandler(updated)
        }
    }

    /// Amp's balance RPC should not redirect. Refusing redirects guarantees the bearer token cannot cross hosts.
    private final class APIRedirectDiagnostics: NSObject, URLSessionTaskDelegate, @unchecked Sendable {
        private let logger: ((String) -> Void)?

        init(logger: ((String) -> Void)?) {
            self.logger = logger
        }

        func urlSession(
            _: URLSession,
            task _: URLSessionTask,
            willPerformHTTPRedirection response: HTTPURLResponse,
            newRequest request: URLRequest,
            completionHandler: @escaping (URLRequest?) -> Void)
        {
            let from = response.url?.absoluteString ?? "unknown"
            let to = request.url?.absoluteString ?? "unknown"
            self.logger?("[amp] API redirect blocked: \(response.statusCode) \(from) -> \(to)")
            completionHandler(nil)
        }
    }

    private struct ResponseInfo {
        let statusCode: Int
        let url: String
    }

    private struct UsageAPIResponse: Decodable {
        let ok: Bool
        let result: Result?
        let error: APIError?

        struct Result: Decodable {
            let displayText: String
        }

        struct APIError: Decodable {
            let code: String?
            let message: String?
        }
    }

    private func logDiagnostics(
        responseInfo: ResponseInfo?,
        diagnostics: RedirectDiagnostics,
        logger: (String) -> Void)
    {
        if let responseInfo {
            logger("[amp] Response: \(responseInfo.statusCode) \(responseInfo.url)")
        }
        if !diagnostics.redirects.isEmpty {
            logger("[amp] Redirects:")
            for entry in diagnostics.redirects {
                logger("[amp]   \(entry)")
            }
        }
    }

    private func cookieNames(from header: String) -> [String] {
        header.split(separator: ";").compactMap { part in
            let trimmed = part.trimmingCharacters(in: .whitespacesAndNewlines)
            guard let idx = trimmed.firstIndex(of: "=") else { return nil }
            let name = trimmed[..<idx].trimmingCharacters(in: .whitespacesAndNewlines)
            return name.isEmpty ? nil : String(name)
        }
    }

    private func sessionCookieHeader(from header: String) -> String? {
        let pairs = CookieHeaderNormalizer.pairs(from: header)
        let sessionPairs = pairs.filter { $0.name == "session" }
        guard !sessionPairs.isEmpty else { return nil }
        return sessionPairs.map { "\($0.name)=\($0.value)" }.joined(separator: "; ")
    }

    static func shouldAttachCookie(to url: URL?) -> Bool {
        guard url?.scheme?.lowercased() == "https" else { return false }
        return self.isAmpHost(url)
    }

    private static func isAmpHost(_ url: URL?) -> Bool {
        guard let host = url?.host?.lowercased() else { return false }
        if host == "ampcode.com" || host == "www.ampcode.com" { return true }
        return host.hasSuffix(".ampcode.com")
    }

    static func isLoginRedirect(_ url: URL) -> Bool {
        guard self.isAmpHost(url) else { return false }
        if url.host?.lowercased() == "auth.ampcode.com" { return true }

        let path = url.path.lowercased()
        let components = path.split(separator: "/").map(String.init)
        if components.contains("login") { return true }
        if components.contains("signin") { return true }
        if components.contains("sign-in") { return true }

        // Amp currently redirects to /auth/sign-in?returnTo=... when session is invalid. Keep this slightly broader
        // than one exact path so we keep working if Amp changes auth routes.
        if components.contains("auth") {
            let query = url.query?.lowercased() ?? ""
            if query.contains("returnto=") { return true }
            if query.contains("redirect=") { return true }
            if query.contains("redirectto=") { return true }
        }

        return false
    }
}