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

public struct KimiK2UsageSnapshot: Sendable {
    public let summary: KimiK2UsageSummary

    public init(summary: KimiK2UsageSummary) {
        self.summary = summary
    }

    public func toUsageSnapshot() -> UsageSnapshot {
        self.summary.toUsageSnapshot()
    }
}

public struct KimiK2UsageSummary: Sendable {
    public let consumed: Double
    public let remaining: Double
    public let averageTokens: Double?
    public let updatedAt: Date

    public init(consumed: Double, remaining: Double, averageTokens: Double?, updatedAt: Date) {
        self.consumed = consumed
        self.remaining = remaining
        self.averageTokens = averageTokens
        self.updatedAt = updatedAt
    }

    public func toUsageSnapshot() -> UsageSnapshot {
        let identity = ProviderIdentitySnapshot(
            providerID: .kimik2,
            accountEmail: nil,
            accountOrganization: nil,
            loginMethod: "Credits: \(UsageFormatter.creditsString(from: self.remaining))")
        return UsageSnapshot(
            primary: nil,
            secondary: nil,
            tertiary: nil,
            providerCost: nil,
            updatedAt: self.updatedAt,
            identity: identity)
    }
}

public enum KimiK2UsageError: LocalizedError, Sendable {
    case missingCredentials
    case invalidCredentials
    case networkError(String)
    case apiError(String)
    case parseFailed(String)

    public var errorDescription: String? {
        switch self {
        case .missingCredentials:
            "Missing Kimi K2 API key."
        case .invalidCredentials:
            "Kimi K2 API key is invalid or expired."
        case let .networkError(message):
            "Kimi K2 network error: \(message)"
        case let .apiError(message):
            "Kimi K2 API error: \(message)"
        case let .parseFailed(message):
            "Failed to parse Kimi K2 response: \(message)"
        }
    }
}

public struct KimiK2UsageFetcher: Sendable {
    private static let log = CodexBarLog.logger(LogCategories.kimiK2Usage)
    private static let creditsURL = URL(string: "https://kimi-k2.ai/api/user/credits")!
    private static let jsonSerializer = JSONSerialization.self
    private static let consumedPaths: [[String]] = [
        ["total_credits_consumed"],
        ["totalCreditsConsumed"],
        ["total_credits_used"],
        ["totalCreditsUsed"],
        ["credits_consumed"],
        ["creditsConsumed"],
        ["consumedCredits"],
        ["usedCredits"],
        ["total"],
        ["usage", "total"],
        ["usage", "consumed"],
    ]

    private static let remainingPaths: [[String]] = [
        ["credits_remaining"],
        ["creditsRemaining"],
        ["remaining_credits"],
        ["remainingCredits"],
        ["available_credits"],
        ["availableCredits"],
        ["credits_left"],
        ["creditsLeft"],
        ["usage", "credits_remaining"],
        ["usage", "remaining"],
    ]

    private static let averageTokenPaths: [[String]] = [
        ["average_tokens_per_request"],
        ["averageTokensPerRequest"],
        ["average_tokens"],
        ["averageTokens"],
        ["avg_tokens"],
        ["avgTokens"],
    ]

    private static let timestampPaths: [[String]] = [
        ["updated_at"],
        ["updatedAt"],
        ["timestamp"],
        ["time"],
        ["last_update"],
        ["lastUpdated"],
    ]

    public static func fetchUsage(
        apiKey: String,
        transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) async throws -> KimiK2UsageSnapshot
    {
        let trimmedKey = apiKey.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !trimmedKey.isEmpty else {
            throw KimiK2UsageError.missingCredentials
        }

        var request = URLRequest(url: self.creditsURL)
        request.httpMethod = "GET"
        request.setValue("Bearer \(trimmedKey)", forHTTPHeaderField: "Authorization")
        request.setValue("application/json", forHTTPHeaderField: "Accept")

        let response = try await transport.response(for: request)
        let data = response.data
        switch response.statusCode {
        case 200:
            break
        case 401:
            throw KimiK2UsageError.invalidCredentials
        default:
            let body = String(data: data, encoding: .utf8) ?? "HTTP \(response.statusCode)"
            Self.log.error("Kimi K2 API returned \(response.statusCode): \(body)")
            throw KimiK2UsageError.apiError(body)
        }

        if let jsonString = String(data: data, encoding: .utf8) {
            Self.log.debug("Kimi K2 API response: \(jsonString)")
        }

        let summary = try Self.parseSummary(data: data, headers: response.response.allHeaderFields)
        return KimiK2UsageSnapshot(summary: summary)
    }

    static func _parseSummaryForTesting(
        _ data: Data,
        headers: [AnyHashable: Any] = [:],
        now: Date = Date()) throws -> KimiK2UsageSummary
    {
        try self.parseSummary(data: data, headers: headers, now: now)
    }

    private static func parseSummary(
        data: Data,
        headers: [AnyHashable: Any],
        now: Date = Date()) throws -> KimiK2UsageSummary
    {
        guard let json = try? jsonSerializer.jsonObject(with: data),
              let dictionary = json as? [String: Any]
        else {
            throw KimiK2UsageError.parseFailed("Root JSON is not an object.")
        }

        let contexts = Self.contexts(from: dictionary)
        let consumed = Self.doubleValue(for: Self.consumedPaths, in: contexts) ?? 0
        let remaining = Self.doubleValue(for: Self.remainingPaths, in: contexts)
            ?? Self.doubleValueFromHeaders(headers: headers, key: "x-credits-remaining")
            ?? 0
        let averageTokens = Self.doubleValue(for: Self.averageTokenPaths, in: contexts)
        let updatedAt = Self.dateValue(for: Self.timestampPaths, in: contexts) ?? now

        return KimiK2UsageSummary(
            consumed: consumed,
            remaining: max(0, remaining),
            averageTokens: averageTokens,
            updatedAt: updatedAt)
    }

    private static func contexts(from dictionary: [String: Any]) -> [[String: Any]] {
        var contexts: [[String: Any]] = [dictionary]
        if let data = dictionary["data"] as? [String: Any] {
            contexts.append(data)
            if let dataUsage = data["usage"] as? [String: Any] {
                contexts.append(dataUsage)
            }
            if let dataCredits = data["credits"] as? [String: Any] {
                contexts.append(dataCredits)
            }
        }
        if let result = dictionary["result"] as? [String: Any] {
            contexts.append(result)
            if let resultUsage = result["usage"] as? [String: Any] {
                contexts.append(resultUsage)
            }
            if let resultCredits = result["credits"] as? [String: Any] {
                contexts.append(resultCredits)
            }
        }
        if let usage = dictionary["usage"] as? [String: Any] {
            contexts.append(usage)
        }
        if let credits = dictionary["credits"] as? [String: Any] {
            contexts.append(credits)
        }
        return contexts
    }

    private static func doubleValue(
        for paths: [[String]],
        in contexts: [[String: Any]]) -> Double?
    {
        for path in paths {
            if let raw = self.value(for: path, in: contexts),
               let value = self.double(from: raw)
            {
                return value
            }
        }
        return nil
    }

    private static func dateValue(
        for paths: [[String]],
        in contexts: [[String: Any]]) -> Date?
    {
        for path in paths {
            if let raw = self.value(for: path, in: contexts) {
                if let date = self.date(from: raw) {
                    return date
                }
            }
        }
        return nil
    }

    private static func value(for path: [String], in contexts: [[String: Any]]) -> Any? {
        for context in contexts {
            var cursor: Any? = context
            for key in path {
                if let dict = cursor as? [String: Any] {
                    cursor = dict[key]
                } else {
                    cursor = nil
                }
            }
            if cursor != nil {
                return cursor
            }
        }
        return nil
    }

    private static func double(from raw: Any) -> Double? {
        let value: Double? = if let raw = raw as? Double {
            raw
        } else if let raw = raw as? Int {
            Double(raw)
        } else if let raw = raw as? String {
            Double(raw)
        } else {
            nil
        }
        guard let value, value.isFinite else { return nil }
        return value
    }

    private static func date(from raw: Any) -> Date? {
        if let value = raw as? Date {
            return value
        }
        if let value = raw as? Double {
            return self.dateFromNumeric(value)
        }
        if let value = raw as? Int {
            return self.dateFromNumeric(Double(value))
        }
        if let value = raw as? String {
            let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines)
            if let numeric = Double(trimmed) {
                return self.dateFromNumeric(numeric)
            }
            let formatter = ISO8601DateFormatter()
            formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
            if let date = formatter.date(from: value) {
                return date
            }
            let fallback = ISO8601DateFormatter()
            fallback.formatOptions = [.withInternetDateTime]
            return fallback.date(from: value)
        }
        return nil
    }

    private static func dateFromNumeric(_ value: Double) -> Date? {
        guard value.isFinite, value > 0 else { return nil }
        let seconds = value >= 1_000_000_000_000 ? value / 1000 : value
        guard seconds.isFinite, seconds <= Date.distantFuture.timeIntervalSince1970 else { return nil }
        return Date(timeIntervalSince1970: seconds)
    }

    private static func doubleValueFromHeaders(headers: [AnyHashable: Any], key: String) -> Double? {
        for (headerKey, value) in headers {
            guard let headerKey = headerKey as? String else { continue }
            if headerKey.lowercased() == key.lowercased() {
                return self.double(from: value)
            }
        }
        return nil
    }
}