import Foundation

enum PiSessionCostCacheIO {
    /// Artifact schema version. Pricing changes are tracked separately by `pricingKey`.
    private static let artifactVersion = 8

    private static func defaultCacheRoot() -> URL {
        let root = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first!
        return root.appendingPathComponent("CodexBar", isDirectory: true)
    }

    static func cacheFileURL(cacheRoot: URL? = nil) -> URL {
        let root = cacheRoot ?? self.defaultCacheRoot()
        return root
            .appendingPathComponent("cost-usage", isDirectory: true)
            .appendingPathComponent("pi-sessions-v\(Self.artifactVersion).json", isDirectory: false)
    }

    static func load(cacheRoot: URL? = nil) -> PiSessionCostCache {
        let url = self.cacheFileURL(cacheRoot: cacheRoot)
        guard let data = try? Data(contentsOf: url),
              let decoded = try? JSONDecoder().decode(PiSessionCostCache.self, from: data),
              decoded.version == Self.artifactVersion
        else {
            return PiSessionCostCache(version: Self.artifactVersion)
        }
        return decoded
    }

    static func save(
        cache: PiSessionCostCache,
        cacheRoot: URL? = nil,
        calendar: Calendar = .current)
    {
        let url = self.cacheFileURL(cacheRoot: cacheRoot)
        let dir = url.deletingLastPathComponent()
        try? FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)

        var cache = cache
        cache.timeZoneIdentifier = calendar.timeZone.identifier
        let tmp = dir.appendingPathComponent(".tmp-\(UUID().uuidString).json", isDirectory: false)
        let data = (try? JSONEncoder().encode(cache)) ?? Data()
        do {
            try data.write(to: tmp, options: [.atomic])
            if FileManager.default.fileExists(atPath: url.path) {
                _ = try FileManager.default.replaceItemAt(url, withItemAt: tmp)
            } else {
                try FileManager.default.moveItem(at: tmp, to: url)
            }
        } catch {
            try? FileManager.default.removeItem(at: tmp)
        }
    }
}

struct PiSessionCostCache: Codable {
    var version: Int
    var lastScanUnixMs: Int64 = 0
    var scanSinceKey: String?
    var scanUntilKey: String?
    var timeZoneIdentifier: String?
    var pricingKey: String?
    var daysByProvider: [String: [String: [String: PiPackedUsage]]] = [:]
    var files: [String: PiSessionFileUsage] = [:]

    init(version: Int = 8) {
        self.version = version
    }
}

struct PiSessionFileUsage: Codable {
    var mtimeUnixMs: Int64
    var size: Int64
    var parsedBytes: Int64
    var sessionID: String?
    var lastModelContext: PiModelContext?
    var contributions: [String: [String: [String: PiPackedUsage]]]
    var unkeyedContributions: [String: [String: [String: PiPackedUsage]]]
    var entryUsages: [String: PiSessionEntryUsage]

    init(
        mtimeUnixMs: Int64,
        size: Int64,
        parsedBytes: Int64,
        sessionID: String? = nil,
        lastModelContext: PiModelContext?,
        contributions: [String: [String: [String: PiPackedUsage]]],
        unkeyedContributions: [String: [String: [String: PiPackedUsage]]] = [:],
        entryUsages: [String: PiSessionEntryUsage] = [:])
    {
        self.mtimeUnixMs = mtimeUnixMs
        self.size = size
        self.parsedBytes = parsedBytes
        self.sessionID = sessionID
        self.lastModelContext = lastModelContext
        self.contributions = contributions
        self.unkeyedContributions = unkeyedContributions
        self.entryUsages = entryUsages
    }
}

struct PiSessionEntryUsage: Codable, Equatable {
    var providerRawValue: String
    var dayKey: String
    var modelName: String
    var usage: PiPackedUsage
}

struct PiModelContext: Codable, Equatable {
    var providerRawValue: String
    var modelName: String
}

struct PiPackedUsage: Codable, Equatable {
    var inputTokens: Int = 0
    var cacheReadTokens: Int = 0
    var cacheWriteTokens: Int = 0
    var outputTokens: Int = 0
    var totalTokens: Int = 0
    var costNanos: Int64 = 0
    var costSampleCount: Int = 0
    var usageSampleCount: Int?

    var isZero: Bool {
        self.inputTokens == 0
            && self.cacheReadTokens == 0
            && self.cacheWriteTokens == 0
            && self.outputTokens == 0
            && self.totalTokens == 0
            && self.costNanos == 0
            && self.costSampleCount == 0
            && (self.usageSampleCount ?? 0) == 0
    }
}
