TokenLens/NativeTokenLens/Sources/TokenLensCore/SQLiteUsageCache.swift

import Foundation
import SQLite3

enum SQLiteUsageCacheError: Error {
    case openFailed(String)
    case prepareFailed(String)
    case executeFailed(String)
    case bindFailed(String)
    case stepFailed(String)
}

final class SQLiteUsageCache {
    private static let transient = unsafeBitCast(-1, to: sqlite3_destructor_type.self)

    private let database: OpaquePointer?
    private var calendar: Calendar = {
        var calendar = Calendar.current
        calendar.firstWeekday = 2
        return calendar
    }()

    init(url: URL) throws {
        try FileManager.default.createDirectory(
            at: url.deletingLastPathComponent(),
            withIntermediateDirectories: true
        )

        var database: OpaquePointer?
        let flags = SQLITE_OPEN_CREATE | SQLITE_OPEN_READWRITE | SQLITE_OPEN_FULLMUTEX
        guard sqlite3_open_v2(url.path, &database, flags, nil) == SQLITE_OK else {
            let message = database.flatMap { String(cString: sqlite3_errmsg($0)) } ?? "无法打开 SQLite"
            sqlite3_close(database)
            throw SQLiteUsageCacheError.openFailed(message)
        }
        self.database = database

        try configure()
        try migrate()
    }

    deinit {
        sqlite3_close(database)
    }

    func isFileCurrent(path: String, size: Int64, modifiedAt: TimeInterval) throws -> Bool {
        let statement = try prepare("SELECT size, modified_at, scanned_offset FROM files WHERE path = ? LIMIT 1")
        defer { sqlite3_finalize(statement) }

        try bind(path, to: statement, at: 1)
        guard sqlite3_step(statement) == SQLITE_ROW else {
            return false
        }

        let cachedSize = sqlite3_column_int64(statement, 0)
        let cachedModifiedAt = sqlite3_column_double(statement, 1)
        let scannedOffset = sqlite3_column_int64(statement, 2)
        return cachedSize == size && scannedOffset == size && abs(cachedModifiedAt - modifiedAt) < 0.001
    }

    func fileState(path: String) throws -> CachedFileState? {
        let statement = try prepare("""
            SELECT size, modified_at, scanned_offset, record_count
            FROM files
            WHERE path = ?
            LIMIT 1
            """)
        defer { sqlite3_finalize(statement) }

        try bind(path, to: statement, at: 1)
        guard sqlite3_step(statement) == SQLITE_ROW else {
            return nil
        }

        return CachedFileState(
            size: sqlite3_column_int64(statement, 0),
            modifiedAt: sqlite3_column_double(statement, 1),
            scannedOffset: sqlite3_column_int64(statement, 2),
            recordCount: Int(sqlite3_column_int64(statement, 3))
        )
    }

    func replaceFileRecords(
        path: String,
        size: Int64,
        modifiedAt: TimeInterval,
        scan: (_ insert: (UsageRecord) throws -> Void) throws -> Int64
    ) throws {
        try execute("BEGIN IMMEDIATE TRANSACTION")
        do {
            try deleteRecords(path: path)
            let insertStatement = try prepare("""
                INSERT INTO records (
                    id, provider, source_path, timestamp, session_id, model,
                    input_tokens, output_tokens, cached_tokens, reasoning_tokens,
                    tool_tokens, total_tokens, day_start, hour_start, month_start
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """)
            defer { sqlite3_finalize(insertStatement) }

            var recordCount = 0
            let scannedOffset = try scan { record in
                recordCount += 1
                try self.insert(record, id: "\(path)#\(recordCount)", statement: insertStatement)
            }
            try upsertFile(
                path: path,
                size: size,
                modifiedAt: modifiedAt,
                scannedOffset: scannedOffset,
                recordCount: recordCount
            )
            try execute("COMMIT")
        } catch {
            try? execute("ROLLBACK")
            throw error
        }
    }

    func appendFileRecords(
        path: String,
        size: Int64,
        modifiedAt: TimeInterval,
        startingRecordCount: Int,
        scan: (_ insert: (UsageRecord) throws -> Void) throws -> Int64
    ) throws {
        try execute("BEGIN IMMEDIATE TRANSACTION")
        do {
            let insertStatement = try prepare("""
                INSERT INTO records (
                    id, provider, source_path, timestamp, session_id, model,
                    input_tokens, output_tokens, cached_tokens, reasoning_tokens,
                    tool_tokens, total_tokens, day_start, hour_start, month_start
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """)
            defer { sqlite3_finalize(insertStatement) }

            var appendedCount = 0
            let scannedOffset = try scan { record in
                appendedCount += 1
                let nextRecordIndex = startingRecordCount + appendedCount
                try self.insert(record, id: "\(path)#\(nextRecordIndex)", statement: insertStatement)
            }
            try upsertFile(
                path: path,
                size: size,
                modifiedAt: modifiedAt,
                scannedOffset: scannedOffset,
                recordCount: startingRecordCount + appendedCount
            )
            try execute("COMMIT")
        } catch {
            try? execute("ROLLBACK")
            throw error
        }
    }

    func pruneMissingFiles(currentPaths: Set<String>) throws {
        let statement = try prepare("SELECT path FROM files")
        defer { sqlite3_finalize(statement) }

        var cachedPaths: [String] = []
        while sqlite3_step(statement) == SQLITE_ROW {
            cachedPaths.append(columnText(statement, at: 0))
        }

        for path in cachedPaths where !currentPaths.contains(path) {
            try deleteRecords(path: path)
            let deleteFile = try prepare("DELETE FROM files WHERE path = ?")
            defer { sqlite3_finalize(deleteFile) }
            try bind(path, to: deleteFile, at: 1)
            try stepDone(deleteFile)
        }
    }

    func makeSummary(now: Date = Date(), calendar: Calendar = Calendar.current) throws -> UsageSummary {
        self.calendar = calendarWithMondayStart(calendar)

        let providerTotals = try providerBuckets()
        let todayProviderTotals = try providerBuckets(
            whereClause: "timestamp >= ? AND timestamp < ?",
            bindings: dayBounds(for: now).bindings
        )
        let dailyDetails = try dailyUsageDetails()
        let dailyTrend = dailyDetails
            .sorted { $0.date < $1.date }
            .suffix(30)
            .map { detail in
                DailyPoint(date: detail.date, values: detail.providerTotals)
            }
        let trendPointsByMode = try Dictionary(uniqueKeysWithValues: TrendMode.allCases.map { mode in
            (mode, try trendPoints(mode: mode, now: now))
        })
        let toolRows = Provider.allCases.compactMap { provider -> ToolRow? in
            guard let bucket = providerTotals[provider], bucket.totalTokens > 0 else { return nil }
            return ToolRow(
                provider: provider,
                inputTokens: bucket.inputTokens,
                outputTokens: bucket.outputTokens,
                cachedTokens: bucket.cachedTokens,
                totalTokens: bucket.totalTokens,
                trendTokens: dailyTrend.map { $0.values[provider, default: 0] }
            )
        }
        .sorted { $0.totalTokens > $1.totalTokens }

        let todayTokens = todayProviderTotals.values.reduce(0) { $0 + $1.totalTokens }
        let monthTokens = try tokenSum(whereClause: "timestamp >= ? AND timestamp < ?", bindings: monthBounds(for: now).bindings)
        let totalTokens = providerTotals.values.reduce(0) { $0 + $1.totalTokens }
        let cachedTokens = providerTotals.values.reduce(0) { $0 + $1.cachedTokens }
        let cacheHitRate = totalTokens > 0 ? Int((Double(cachedTokens) / Double(totalTokens) * 100).rounded()) : 0
        let cards = SummaryCards(
            todayTokens: todayTokens,
            monthTokens: monthTokens,
            totalTokens: totalTokens,
            cacheHitRate: cacheHitRate,
            sessionCount: try sessionCount()
        )

        return UsageSummary(
            generatedAt: now,
            records: [],
            cards: cards,
            providerTotals: providerTotals,
            todayProviderTotals: todayProviderTotals,
            dailyTrend: Array(dailyTrend),
            trendPointsByMode: trendPointsByMode,
            dailyDetails: dailyDetails.sorted { $0.date > $1.date },
            toolRows: toolRows,
            insights: UsageSummary.makeInsights(
                providerTotals: providerTotals,
                todayTokens: todayTokens,
                cacheHitRate: cacheHitRate
            )
        )
    }

    private func configure() throws {
        try execute("PRAGMA journal_mode = WAL")
        try execute("PRAGMA synchronous = NORMAL")
        try execute("PRAGMA temp_store = MEMORY")
        try execute("PRAGMA busy_timeout = 3000")
    }

    private func migrate() throws {
        if try recordsTableNeedsRebuild() {
            try execute("DROP TABLE IF EXISTS records")
            try execute("DROP TABLE IF EXISTS files")
        }
        try execute("""
            CREATE TABLE IF NOT EXISTS files (
                path TEXT PRIMARY KEY,
                size INTEGER NOT NULL,
                modified_at REAL NOT NULL,
                scanned_offset INTEGER NOT NULL DEFAULT 0,
                record_count INTEGER NOT NULL DEFAULT 0,
                scanned_at REAL NOT NULL
            )
            """)
        try addColumnIfMissing(
            table: "files",
            column: "scanned_offset",
            definition: "scanned_offset INTEGER NOT NULL DEFAULT 0"
        )
        try addColumnIfMissing(
            table: "files",
            column: "record_count",
            definition: "record_count INTEGER NOT NULL DEFAULT 0"
        )
        try execute("""
            CREATE TABLE IF NOT EXISTS records (
                id TEXT PRIMARY KEY,
                provider TEXT NOT NULL,
                source_path TEXT NOT NULL,
                timestamp REAL NOT NULL,
                session_id TEXT NOT NULL,
                model TEXT NOT NULL,
                input_tokens INTEGER NOT NULL,
                output_tokens INTEGER NOT NULL,
                cached_tokens INTEGER NOT NULL,
                reasoning_tokens INTEGER NOT NULL,
                tool_tokens INTEGER NOT NULL,
                total_tokens INTEGER NOT NULL,
                day_start REAL NOT NULL,
                hour_start REAL NOT NULL,
                month_start REAL NOT NULL
            )
            """)
        try execute("CREATE INDEX IF NOT EXISTS records_provider_idx ON records(provider)")
        try execute("CREATE INDEX IF NOT EXISTS records_source_path_idx ON records(source_path)")
        try execute("CREATE INDEX IF NOT EXISTS records_timestamp_idx ON records(timestamp)")
        try execute("CREATE INDEX IF NOT EXISTS records_day_provider_idx ON records(day_start, provider)")
        try execute("CREATE INDEX IF NOT EXISTS records_hour_provider_idx ON records(hour_start, provider)")
        try execute("CREATE INDEX IF NOT EXISTS records_month_provider_idx ON records(month_start, provider)")
    }

    private func deleteRecords(path: String) throws {
        let statement = try prepare("DELETE FROM records WHERE source_path = ?")
        defer { sqlite3_finalize(statement) }
        try bind(path, to: statement, at: 1)
        try stepDone(statement)
    }

    private func upsertFile(
        path: String,
        size: Int64,
        modifiedAt: TimeInterval,
        scannedOffset: Int64,
        recordCount: Int
    ) throws {
        let statement = try prepare("""
            INSERT INTO files (path, size, modified_at, scanned_offset, record_count, scanned_at)
            VALUES (?, ?, ?, ?, ?, ?)
            ON CONFLICT(path) DO UPDATE SET
                size = excluded.size,
                modified_at = excluded.modified_at,
                scanned_offset = excluded.scanned_offset,
                record_count = excluded.record_count,
                scanned_at = excluded.scanned_at
            """)
        defer { sqlite3_finalize(statement) }

        try bind(path, to: statement, at: 1)
        sqlite3_bind_int64(statement, 2, size)
        sqlite3_bind_double(statement, 3, modifiedAt)
        sqlite3_bind_int64(statement, 4, scannedOffset)
        sqlite3_bind_int64(statement, 5, Int64(recordCount))
        sqlite3_bind_double(statement, 6, Date().timeIntervalSince1970)
        try stepDone(statement)
    }

    private func insert(_ record: UsageRecord, id: String, statement: OpaquePointer?) throws {
        sqlite3_reset(statement)
        sqlite3_clear_bindings(statement)

        try bind(id, to: statement, at: 1)
        try bind(record.provider.rawValue, to: statement, at: 2)
        try bind(record.sourcePath, to: statement, at: 3)
        sqlite3_bind_double(statement, 4, record.timestamp.timeIntervalSince1970)
        try bind(record.sessionID, to: statement, at: 5)
        try bind(record.model, to: statement, at: 6)
        sqlite3_bind_int64(statement, 7, Int64(record.inputTokens))
        sqlite3_bind_int64(statement, 8, Int64(record.outputTokens))
        sqlite3_bind_int64(statement, 9, Int64(record.cachedTokens))
        sqlite3_bind_int64(statement, 10, Int64(record.reasoningTokens))
        sqlite3_bind_int64(statement, 11, Int64(record.toolTokens))
        sqlite3_bind_int64(statement, 12, Int64(record.totalTokens))
        sqlite3_bind_double(statement, 13, calendar.startOfDay(for: record.timestamp).timeIntervalSince1970)
        sqlite3_bind_double(statement, 14, hourStart(for: record.timestamp).timeIntervalSince1970)
        sqlite3_bind_double(statement, 15, monthStart(for: record.timestamp).timeIntervalSince1970)
        try stepDone(statement)
    }

    private func providerBuckets(
        whereClause: String? = nil,
        bindings: [TimeInterval] = []
    ) throws -> [Provider: TokenBucket] {
        var buckets = Dictionary(uniqueKeysWithValues: Provider.allCases.map {
            ($0, TokenBucket(provider: $0))
        })
        var sql = """
            SELECT provider,
                   SUM(input_tokens), SUM(output_tokens), SUM(cached_tokens),
                   SUM(reasoning_tokens), SUM(tool_tokens), SUM(total_tokens), COUNT(*)
            FROM records
            """
        if let whereClause {
            sql += " WHERE \(whereClause)"
        }
        sql += " GROUP BY provider"

        let statement = try prepare(sql)
        defer { sqlite3_finalize(statement) }
        try bindIntervals(bindings, to: statement)

        while sqlite3_step(statement) == SQLITE_ROW {
            guard let provider = Provider(rawValue: columnText(statement, at: 0)) else { continue }
            var bucket = TokenBucket(provider: provider)
            bucket.inputTokens = Int(sqlite3_column_int64(statement, 1))
            bucket.outputTokens = Int(sqlite3_column_int64(statement, 2))
            bucket.cachedTokens = Int(sqlite3_column_int64(statement, 3))
            bucket.reasoningTokens = Int(sqlite3_column_int64(statement, 4))
            bucket.toolTokens = Int(sqlite3_column_int64(statement, 5))
            bucket.totalTokens = Int(sqlite3_column_int64(statement, 6))
            bucket.count = Int(sqlite3_column_int64(statement, 7))
            buckets[provider] = bucket
        }

        return buckets
    }

    private func dailyUsageDetails() throws -> [DailyUsageDetail] {
        let statement = try prepare("""
            SELECT day_start, provider,
                   SUM(input_tokens), SUM(output_tokens), SUM(cached_tokens), SUM(total_tokens)
            FROM records
            GROUP BY day_start, provider
            ORDER BY day_start DESC
            """)
        defer { sqlite3_finalize(statement) }

        var details: [Date: DailyUsageDetail] = [:]
        while sqlite3_step(statement) == SQLITE_ROW {
            guard let provider = Provider(rawValue: columnText(statement, at: 1)) else { continue }
            let date = Date(timeIntervalSince1970: sqlite3_column_double(statement, 0))
            if details[date] == nil {
                details[date] = DailyUsageDetail(date: date)
            }
            details[date]?.add(
                provider: provider,
                inputTokens: Int(sqlite3_column_int64(statement, 2)),
                outputTokens: Int(sqlite3_column_int64(statement, 3)),
                cachedTokens: Int(sqlite3_column_int64(statement, 4)),
                totalTokens: Int(sqlite3_column_int64(statement, 5))
            )
        }

        return details.values.sorted { $0.date > $1.date }
    }

    private func trendPoints(mode: TrendMode, now: Date) throws -> [DailyPoint] {
        let reference = mode.periodStart(for: now, calendar: calendar)
        let start = mode.startDate(reference: reference, calendar: calendar)
        let column = mode.periodColumn
        let statement = try prepare("""
            SELECT \(column), provider, SUM(total_tokens)
            FROM records
            WHERE \(column) >= ? AND \(column) <= ?
            GROUP BY \(column), provider
            """)
        defer { sqlite3_finalize(statement) }
        sqlite3_bind_double(statement, 1, start.timeIntervalSince1970)
        sqlite3_bind_double(statement, 2, reference.timeIntervalSince1970)

        var buckets: [Date: [Provider: Int]] = [:]
        while sqlite3_step(statement) == SQLITE_ROW {
            guard let provider = Provider(rawValue: columnText(statement, at: 1)) else { continue }
            let date = Date(timeIntervalSince1970: sqlite3_column_double(statement, 0))
            buckets[date, default: Provider.zeroTokenMap][provider, default: 0] = Int(sqlite3_column_int64(statement, 2))
        }

        var points: [DailyPoint] = []
        var cursor = start
        while cursor <= reference {
            points.append(DailyPoint(date: cursor, values: buckets[cursor, default: Provider.zeroTokenMap]))
            guard let next = calendar.date(byAdding: mode.calendarComponent, value: 1, to: cursor) else { break }
            cursor = next
        }
        return points
    }

    private func tokenSum(whereClause: String, bindings: [TimeInterval]) throws -> Int {
        let statement = try prepare("SELECT COALESCE(SUM(total_tokens), 0) FROM records WHERE \(whereClause)")
        defer { sqlite3_finalize(statement) }
        try bindIntervals(bindings, to: statement)
        guard sqlite3_step(statement) == SQLITE_ROW else { return 0 }
        return Int(sqlite3_column_int64(statement, 0))
    }

    private func sessionCount() throws -> Int {
        let statement = try prepare("""
            SELECT COUNT(DISTINCT CASE WHEN session_id = '' THEN source_path ELSE session_id END)
            FROM records
            """)
        defer { sqlite3_finalize(statement) }
        guard sqlite3_step(statement) == SQLITE_ROW else { return 0 }
        return Int(sqlite3_column_int64(statement, 0))
    }

    private func recordsTableNeedsRebuild() throws -> Bool {
        let columns = try tableColumns("records")
        return !columns.isEmpty && !columns.contains("day_start")
    }

    private func addColumnIfMissing(table: String, column: String, definition: String) throws {
        let columns = try tableColumns(table)
        guard !columns.isEmpty, !columns.contains(column) else { return }
        try execute("ALTER TABLE \(table) ADD COLUMN \(definition)")
    }

    private func tableColumns(_ table: String) throws -> Set<String> {
        let statement = try prepare("PRAGMA table_info(\(table))")
        defer { sqlite3_finalize(statement) }

        var columns = Set<String>()
        while sqlite3_step(statement) == SQLITE_ROW {
            columns.insert(columnText(statement, at: 1))
        }
        return columns
    }

    private func bindIntervals(_ values: [TimeInterval], to statement: OpaquePointer?) throws {
        for (index, value) in values.enumerated() {
            sqlite3_bind_double(statement, Int32(index + 1), value)
        }
    }

    private func dayBounds(for date: Date) -> QueryBounds {
        let start = calendar.startOfDay(for: date)
        let end = calendar.date(byAdding: .day, value: 1, to: start) ?? date
        return QueryBounds(start: start, end: end)
    }

    private func monthBounds(for date: Date) -> QueryBounds {
        let start = monthStart(for: date)
        let end = calendar.date(byAdding: .month, value: 1, to: start) ?? date
        return QueryBounds(start: start, end: end)
    }

    private func hourStart(for date: Date) -> Date {
        let components = calendar.dateComponents([.year, .month, .day, .hour], from: date)
        return calendar.date(from: components) ?? calendar.startOfDay(for: date)
    }

    private func monthStart(for date: Date) -> Date {
        let components = calendar.dateComponents([.year, .month], from: date)
        return calendar.date(from: components) ?? calendar.startOfDay(for: date)
    }

    private func calendarWithMondayStart(_ calendar: Calendar) -> Calendar {
        var copy = calendar
        copy.firstWeekday = 2
        return copy
    }

    private func prepare(_ sql: String) throws -> OpaquePointer? {
        var statement: OpaquePointer?
        guard sqlite3_prepare_v2(database, sql, -1, &statement, nil) == SQLITE_OK else {
            throw SQLiteUsageCacheError.prepareFailed(errorMessage)
        }
        return statement
    }

    private func execute(_ sql: String) throws {
        guard sqlite3_exec(database, sql, nil, nil, nil) == SQLITE_OK else {
            throw SQLiteUsageCacheError.executeFailed(errorMessage)
        }
    }

    private func bind(_ text: String, to statement: OpaquePointer?, at index: Int32) throws {
        let result = text.withCString {
            sqlite3_bind_text(statement, index, $0, -1, Self.transient)
        }
        guard result == SQLITE_OK else {
            throw SQLiteUsageCacheError.bindFailed(errorMessage)
        }
    }

    private func stepDone(_ statement: OpaquePointer?) throws {
        guard sqlite3_step(statement) == SQLITE_DONE else {
            throw SQLiteUsageCacheError.stepFailed(errorMessage)
        }
    }

    private func columnText(_ statement: OpaquePointer?, at index: Int32) -> String {
        guard let text = sqlite3_column_text(statement, index) else {
            return ""
        }
        return String(cString: UnsafeRawPointer(text).assumingMemoryBound(to: CChar.self))
    }

    private var errorMessage: String {
        database.flatMap { String(cString: sqlite3_errmsg($0)) } ?? "SQLite 操作失败"
    }

}

struct CachedFileState {
    let size: Int64
    let modifiedAt: TimeInterval
    let scannedOffset: Int64
    let recordCount: Int
}

private struct QueryBounds {
    let start: Date
    let end: Date

    var bindings: [TimeInterval] {
        [start.timeIntervalSince1970, end.timeIntervalSince1970]
    }
}