import CUsage
import Foundation

/// One quota window, matching the contract in packages/ai/docs/usage-feeds.md.
public struct UsageLimit: Codable, Sendable, Equatable {
	/// Stable machine key within the provider; `label` is display text.
	public var id: String
	public var label: String
	public var usedPercent: Double?
	public var windowLabel: String?
	public var windowSeconds: Int64?
	public var resetsAt: Int64?
	/// The model family or feature this limit is restricted to; nil counts all usage.
	public var scope: String?
	public var status: String
	public var detail: String?
}

/// What a provider's feed reports for one account.
public struct UsageReport: Codable, Sendable, Equatable {
	public var provider: String?
	public var providerName: String?
	public var accountId: String?
	public var email: String?
	public var planType: String?
	public var limits: [UsageLimit]
	public var balance: String?
	public var notes: [String]?
	public var error: String?
}

/// A provider with a feed, as the core describes it.
public struct UsageProviderInfo: Codable, Sendable, Equatable {
	public var provider: String
	public var providerName: String
	public var credentialKinds: [String]
}

/// The request the core wants sent. The host does the sending.
///
/// Headers arrive as ordered pairs across the C ABI, because HTTP allows a name to repeat and a
/// dictionary cannot. `headers` is the convenience view for the common case.
public struct UsageRequest: Codable, Sendable, Equatable {
	public var kind: String = "http"

	// HTTP request fields
	public var method: String = "GET"
	public var url: String = ""
	public var headerPairs: [[String]]?
	public var body: String?

	// Command request fields
	public var command: String?
	public var args: [String]?
	public var stdin: String?
	public var timeoutMs: Int?

	public var headers: [String: String] {
		Dictionary(
			(headerPairs ?? []).compactMap { pair in pair.count == 2 ? (pair[0], pair[1]) : nil },
			uniquingKeysWith: { _, last in last }
		)
	}

	enum CodingKeys: String, CodingKey {
		case kind, method, url, body, command, args, stdin, timeoutMs
		case headerPairs = "headers"
	}

	public init(
		kind: String = "http",
		method: String = "GET",
		url: String = "",
		headerPairs: [[String]]? = nil,
		body: String? = nil,
		command: String? = nil,
		args: [String]? = nil,
		stdin: String? = nil,
		timeoutMs: Int? = nil
	) {
		self.kind = kind
		self.method = method
		self.url = url
		self.headerPairs = headerPairs
		self.body = body
		self.command = command
		self.args = args
		self.stdin = stdin
		self.timeoutMs = timeoutMs
	}

	public init(from decoder: Decoder) throws {
		let container = try decoder.container(keyedBy: CodingKeys.self)
		self.kind = try container.decodeIfPresent(String.self, forKey: .kind) ?? "http"
		self.method = try container.decodeIfPresent(String.self, forKey: .method) ?? "GET"
		self.url = try container.decodeIfPresent(String.self, forKey: .url) ?? ""
		self.headerPairs = try container.decodeIfPresent([[String]].self, forKey: .headerPairs)
		self.body = try container.decodeIfPresent(String.self, forKey: .body)
		self.command = try container.decodeIfPresent(String.self, forKey: .command)
		self.args = try container.decodeIfPresent([String].self, forKey: .args)
		self.stdin = try container.decodeIfPresent(String.self, forKey: .stdin)
		self.timeoutMs = try container.decodeIfPresent(Int.self, forKey: .timeoutMs)
	}
}

/// One response handed back to the core.
public struct UsageResponse: Codable, Sendable, Equatable {
	public var status: Int
	public var body: String

	public init(status: Int, body: String) {
		self.status = status
		self.body = body
	}
}

/// Either the next requests to send, or the finished report.
public enum UsageStep: Sendable {
	case requests([UsageRequest])
	case report(UsageReport)

	/// A provider that could not report says so in its report rather than by throwing.
	public var reportValue: UsageReport? {
		if case .report(let report) = self { return report }
		return nil
	}
}

/// Opaque cross-round state managed by the host.
public final class UsageState: @unchecked Sendable {
	let pointer: OpaquePointer?

	public init() {
		self.pointer = usage_state_new()
	}

	deinit {
		if let pointer {
			usage_state_free(pointer)
		}
	}
}

public enum UsageCoreError: Error, CustomStringConvertible {
	case missingLibraryResult
	case invalidJSON(String)
	case providerRefused(String)
	case incompatibleAbiVersion(expected: UInt32, actual: UInt32)

	public var description: String {
		switch self {
		case .missingLibraryResult: return "the usage core returned nothing"
		case .invalidJSON(let text): return "unexpected JSON from the usage core: \(text)"
		case .providerRefused(let message): return message
		case .incompatibleAbiVersion(let expected, let actual):
			return "incompatible usage core ABI version: expected \(expected), got \(actual)"
		}
	}
}

public enum UsageCore {
	public static let expectedAbiVersion: UInt32 = 4
	public static var abiVersion: UInt32 { usage_abi_version() }

	public static func checkAbiVersion() throws {
		let version = abiVersion
		guard version == expectedAbiVersion else {
			throw UsageCoreError.incompatibleAbiVersion(expected: expectedAbiVersion, actual: version)
		}
	}

	public static func providers() throws -> [UsageProviderInfo] {
		try checkAbiVersion()
		let json = try consume(usage_providers())
		return try JSONDecoder().decode([UsageProviderInfo].self, from: Data(json.utf8))
	}

	/// One step of the host loop: hand the responses you already have (none on the first call) and
	/// get back either the requests to send or the report.
	public static func advance(
		state: UsageState? = nil,
		provider: String,
		credentialJSON: String,
		responses: [UsageResponse],
		nowMs: Int64
	) throws -> UsageStep {
		try checkAbiVersion()
		let responsesJSON = String(data: try JSONEncoder().encode(responses), encoding: .utf8) ?? "[]"

		let raw = try consume(
			provider.withCString { providerPtr in
				credentialJSON.withCString { credentialPtr in
					responsesJSON.withCString { responsesPtr in
						usage_advance(
							state?.pointer,
							providerPtr, strlen(providerPtr),
							credentialPtr, strlen(credentialPtr),
							responsesPtr, strlen(responsesPtr),
							nowMs
						)
					}
				}
			}
		)

		guard let data = raw.data(using: .utf8), let object = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
			throw UsageCoreError.invalidJSON(raw)
		}
		if object["requests"] != nil {
			return .requests(try JSONDecoder().decode(UsageRequestsEnvelope.self, from: data).requests)
		}
		if let error = object["error"] as? String {
			throw UsageCoreError.providerRefused(error)
		}
		return .report(try JSONDecoder().decode(UsageReportEnvelope.self, from: data).report)
	}

	/// Account id and email from an OAuth access token, without a request.
	public static func credentialsFromAccessToken(_ jwt: String) throws -> [String: String] {
		let json = try consume(
			jwt.withCString { pointer in
				usage_credentials_from_access_token(pointer, strlen(pointer))
			}
		)
		guard let data = json.data(using: .utf8) else { throw UsageCoreError.invalidJSON(json) }
		return (try? JSONDecoder().decode([String: String].self, from: data)) ?? [:]
	}

	private struct UsageRequestsEnvelope: Codable { var requests: [UsageRequest] }
	private struct UsageReportEnvelope: Codable { var report: UsageReport }

	/// Every call hands back a string the caller owns. NULL means allocation failed, which for a
	/// quota read is a failure like any other, so it becomes an error instead of a crash.
	private static func consume(_ pointer: UnsafeMutablePointer<CChar>?) throws -> String {
		guard let pointer else { throw UsageCoreError.missingLibraryResult }
		defer { usage_free(pointer) }
		return String(cString: pointer)
	}
}