import Foundation
import SQLite3

public actor TokenUsageScanner {
    public let homeDirectory: URL
    public let cacheURL: URL
    private let chunkSize = 256 * 1024
    private var lastSignature: FileSignature?
    private var lastSummary: UsageSummary?

    public init(
        homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser,
        cacheURL: URL? = nil
    ) {
        self.homeDirectory = homeDirectory
        self.cacheURL = cacheURL ?? homeDirectory
            .appending(path: "Library/Application Support/TokenLens")
            .appending(path: "usage-cache.sqlite")
    }

    public func scan(now: Date = Date()) async -> UsageSummary {
        defer { MemoryPressure.relieve() }

        if let summary = scanWithCache(now: now) {
            return summary
        }

        var accumulator = UsageSummaryAccumulator(now: now)
        scanClaude(into: &accumulator)
        scanCodex(into: &accumulator)
        scanGemini(into: &accumulator)
        scanHermes(into: &accumulator)
        scanOpenCode(into: &accumulator)
        return accumulator.makeSummary()
    }

    private func scanWithCache(now: Date) -> UsageSummary? {
        do {
            let snapshots = try logFileSnapshots()
            let signature = FileSignature(snapshots: snapshots)
            if signature == lastSignature,
               let lastSummary,
               isSameDisplayWindow(lastSummary.generatedAt, now) {
                let updated = lastSummary.updatingGeneratedAt(now)
                self.lastSummary = updated
                return updated
            }

            let cache = try SQLiteUsageCache(url: cacheURL)
            let currentPaths = try refreshCache(cache, snapshots: snapshots)
            try cache.pruneMissingFiles(currentPaths: currentPaths)
            let summary = try cache.makeSummary(now: now)
            lastSignature = signature
            lastSummary = summary
            return summary
        } catch {
            return nil
        }
    }

    private func refreshCache(_ cache: SQLiteUsageCache, snapshots: [LogFileSnapshot]) throws -> Set<String> {
        var currentPaths = Set<String>()
        for snapshot in snapshots {
            currentPaths.insert(snapshot.path)
            let state = try cache.fileState(path: snapshot.path)
            let isCurrent = try cache.isFileCurrent(
                path: snapshot.path,
                size: snapshot.size,
                modifiedAt: snapshot.modifiedAt
            )
            if isCurrent && (snapshot.scanner.supportsZeroRecordCache || (state?.recordCount ?? 0) > 0) {
                continue
            }

            if let state,
               snapshot.size > state.size,
               state.size == state.scannedOffset,
               state.scannedOffset > 0,
               snapshot.scanner.supportsIncrementalScan {
                try cache.appendFileRecords(
                    path: snapshot.path,
                    size: snapshot.size,
                    modifiedAt: snapshot.modifiedAt,
                    startingRecordCount: state.recordCount
                ) { insert in
                    try scanFile(
                        snapshot.url,
                        startingAt: UInt64(state.scannedOffset),
                        scanner: snapshot.scanner,
                        consume: insert
                    ).scannedOffset
                }
            } else {
                try cache.replaceFileRecords(
                    path: snapshot.path,
                    size: snapshot.size,
                    modifiedAt: snapshot.modifiedAt
                ) { insert in
                    try scanFile(snapshot.url, scanner: snapshot.scanner, consume: insert).scannedOffset
                }
            }
        }
        return currentPaths
    }

    private func scanClaude(into accumulator: inout UsageSummaryAccumulator) {
        scanJSONLFiles(
            root: homeDirectory.appending(path: ".claude/projects"),
            include: { $0.pathExtension == "jsonl" },
            parse: TokenUsageParser.parseClaudeLine,
            into: &accumulator
        )
    }

    private func scanCodex(into accumulator: inout UsageSummaryAccumulator) {
        scanJSONLFiles(
            root: homeDirectory.appending(path: ".codex/sessions"),
            include: { $0.pathExtension == "jsonl" },
            parse: TokenUsageParser.parseCodexLine,
            into: &accumulator
        )
    }

    private func scanGemini(into accumulator: inout UsageSummaryAccumulator) {
        scanJSONLFiles(
            root: homeDirectory.appending(path: ".gemini"),
            include: { url in
                url.pathExtension == "jsonl" && url.pathComponents.contains("chats")
            },
            parse: TokenUsageParser.parseGeminiLine,
            into: &accumulator
        )
    }

    private func scanHermes(into accumulator: inout UsageSummaryAccumulator) {
        for file in files(root: homeDirectory.appending(path: ".hermes"), include: { $0.lastPathComponent == "state.db" }) {
            _ = try? scanHermesStateDB(file) { record in
                accumulator.add(record)
            }
        }
    }

    private func scanOpenCode(into accumulator: inout UsageSummaryAccumulator) {
        for file in files(root: homeDirectory.appending(path: ".local/share/opencode"), include: { $0.lastPathComponent == "opencode.db" }) {
            _ = try? scanOpenCodeDB(file) { record in
                accumulator.add(record)
            }
        }
    }

    private func scanJSONLFiles(
        root: URL,
        include: (URL) -> Bool,
        parse: (String, String) -> UsageRecord?,
        into accumulator: inout UsageSummaryAccumulator
    ) {
        for file in files(root: root, include: include) {
            _ = try? scanJSONLFile(file, parse: parse) { record in
                accumulator.add(record)
            }
        }
    }

    private func scanFile(
        _ file: URL,
        startingAt offset: UInt64 = 0,
        scanner: SourceScanner,
        consume: (UsageRecord) throws -> Void
    ) throws -> ScanProgress {
        switch scanner {
        case let .jsonl(parse):
            return try scanJSONLFile(file, startingAt: offset, parse: parse, consume: consume)
        case .hermesStateDB:
            return try scanHermesStateDB(file, consume: consume)
        case .openCodeDB:
            return try scanOpenCodeDB(file, consume: consume)
        }
    }

    private func scanJSONLFile(
        _ file: URL,
        startingAt offset: UInt64 = 0,
        parse: (String, String) -> UsageRecord?,
        consume: (UsageRecord) throws -> Void
    ) throws -> ScanProgress {
        guard let handle = try? FileHandle(forReadingFrom: file) else {
            return ScanProgress(scannedOffset: Int64(offset), recordCount: 0)
        }
        defer { try? handle.close() }
        try handle.seek(toOffset: offset)

        var pendingLine = Data()
        var scannedOffset = Int64(offset)
        var recordCount = 0

        while true {
            let chunk = (try? handle.read(upToCount: chunkSize)) ?? Data()
            if chunk.isEmpty { break }

            var lineStart = chunk.startIndex
            for index in chunk.indices where chunk[index] == 0x0A {
                let lineData = chunk[lineStart..<index]
                if pendingLine.isEmpty {
                    scannedOffset += Int64(lineData.count + 1)
                    try autoreleasepool {
                        if try appendRecord(from: lineData, sourcePath: file.path, parse: parse, consume: consume) {
                            recordCount += 1
                        }
                    }
                } else {
                    pendingLine.append(contentsOf: lineData)
                    scannedOffset += Int64(pendingLine.count + 1)
                    try autoreleasepool {
                        if try appendRecord(from: pendingLine, sourcePath: file.path, parse: parse, consume: consume) {
                            recordCount += 1
                        }
                    }
                    pendingLine.removeAll(keepingCapacity: true)
                }
                lineStart = chunk.index(after: index)
            }

            if lineStart < chunk.endIndex {
                pendingLine.append(contentsOf: chunk[lineStart..<chunk.endIndex])
            }
        }

        if !pendingLine.isEmpty {
            try autoreleasepool {
                if try appendRecord(from: pendingLine, sourcePath: file.path, parse: parse, consume: consume) {
                    recordCount += 1
                    scannedOffset += Int64(pendingLine.count)
                } else if isCompleteJSONLine(pendingLine) {
                    scannedOffset += Int64(pendingLine.count)
                }
            }
        }

        return ScanProgress(scannedOffset: scannedOffset, recordCount: recordCount)
    }

    private func appendRecord<LineData: Collection>(
        from lineData: LineData,
        sourcePath: String,
        parse: (String, String) -> UsageRecord?,
        consume: (UsageRecord) throws -> Void
    ) throws -> Bool where LineData.Element == UInt8 {
        guard !lineData.isEmpty,
              let line = String(bytes: lineData, encoding: .utf8),
              let record = parse(line, sourcePath)
        else {
            return false
        }
        try consume(record)
        return true
    }

    private func isCompleteJSONLine(_ lineData: Data) -> Bool {
        guard !lineData.isEmpty else { return true }
        return (try? JSONSerialization.jsonObject(with: lineData)) != nil
    }

    private func scanHermesStateDB(
        _ file: URL,
        consume: (UsageRecord) throws -> Void
    ) throws -> ScanProgress {
        let sql = """
            SELECT id, started_at, COALESCE(model, ''),
                   input_tokens, output_tokens, cache_read_tokens,
                   cache_write_tokens, reasoning_tokens
            FROM sessions
            WHERE input_tokens + output_tokens + cache_read_tokens + cache_write_tokens + reasoning_tokens > 0
            """
        return try scanSQLiteRows(file, sql: sql) { statement in
            UsageRecord(
                provider: .hermes,
                source: "state.db/sessions",
                sourcePath: file.path,
                timestamp: Date(timeIntervalSince1970: sqlite3_column_double(statement, 1)),
                sessionID: textColumn(statement, 0),
                model: textColumn(statement, 2).isEmpty ? "Hermes" : textColumn(statement, 2),
                inputTokens: intColumn(statement, 3),
                outputTokens: intColumn(statement, 4),
                cachedTokens: intColumn(statement, 5) + intColumn(statement, 6),
                reasoningTokens: intColumn(statement, 7),
                toolTokens: 0
            )
        } consume: {
            try consume($0)
        }
    }

    private func scanOpenCodeDB(
        _ file: URL,
        consume: (UsageRecord) throws -> Void
    ) throws -> ScanProgress {
        let sql = """
            SELECT id, COALESCE(NULLIF(time_updated, 0), time_created), COALESCE(model, ''),
                   tokens_input, tokens_output, tokens_reasoning,
                   tokens_cache_read, tokens_cache_write
            FROM session
            WHERE tokens_input + tokens_output + tokens_reasoning + tokens_cache_read + tokens_cache_write > 0
            """
        return try scanSQLiteRows(file, sql: sql) { statement in
            UsageRecord(
                provider: .opencode,
                source: "opencode.db/session",
                sourcePath: file.path,
                timestamp: Date(timeIntervalSince1970: sqlite3_column_double(statement, 1) / 1000),
                sessionID: textColumn(statement, 0),
                model: openCodeModelName(textColumn(statement, 2)),
                inputTokens: intColumn(statement, 3),
                outputTokens: intColumn(statement, 4),
                cachedTokens: intColumn(statement, 6) + intColumn(statement, 7),
                reasoningTokens: intColumn(statement, 5),
                toolTokens: 0
            )
        } consume: {
            try consume($0)
        }
    }

    private func scanSQLiteRows(
        _ file: URL,
        sql: String,
        makeRecord: (OpaquePointer?) -> UsageRecord,
        consume: (UsageRecord) throws -> Void
    ) throws -> ScanProgress {
        var database: OpaquePointer?
        let flags = SQLITE_OPEN_READWRITE | SQLITE_OPEN_FULLMUTEX
        guard sqlite3_open_v2(file.path, &database, flags, nil) == SQLITE_OK else {
            sqlite3_close(database)
            return ScanProgress(scannedOffset: fileSize(file), recordCount: 0)
        }
        defer { sqlite3_close(database) }
        sqlite3_busy_timeout(database, 1000)

        var statement: OpaquePointer?
        guard sqlite3_prepare_v2(database, sql, -1, &statement, nil) == SQLITE_OK else {
            sqlite3_finalize(statement)
            return ScanProgress(scannedOffset: fileSize(file), recordCount: 0)
        }
        defer { sqlite3_finalize(statement) }

        var recordCount = 0
        while sqlite3_step(statement) == SQLITE_ROW {
            try consume(makeRecord(statement))
            recordCount += 1
        }
        return ScanProgress(scannedOffset: fileSize(file), recordCount: recordCount)
    }

    private func fileSize(_ file: URL) -> Int64 {
        ((try? file.resourceValues(forKeys: [.fileSizeKey]).fileSize).flatMap { Int64($0) }) ?? 0
    }

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

    private func intColumn(_ statement: OpaquePointer?, _ index: Int32) -> Int {
        Int(sqlite3_column_int64(statement, index))
    }

    private func openCodeModelName(_ rawModel: String) -> String {
        guard !rawModel.isEmpty,
              let data = rawModel.data(using: .utf8),
              let object = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any]
        else {
            return rawModel.isEmpty ? "OpenCode" : rawModel
        }
        let provider = object["providerID"] as? String
        let model = object["id"] as? String
        var parts: [String] = []
        if let provider, !provider.isEmpty {
            parts.append(provider)
        }
        if let model, !model.isEmpty {
            parts.append(model)
        }
        return parts.isEmpty ? "OpenCode" : parts.joined(separator: "/")
    }

    private func files(root: URL, include: (URL) -> Bool) -> [URL] {
        guard let enumerator = FileManager.default.enumerator(
            at: root,
            includingPropertiesForKeys: [.isRegularFileKey],
            options: [.skipsHiddenFiles]
        ) else {
            return []
        }

        var result: [URL] = []
        for case let file as URL in enumerator where include(file) {
            result.append(file)
        }
        return result
    }

    private var logSources: [LogSource] {
        [
            LogSource(
                root: homeDirectory.appending(path: ".claude/projects"),
                include: { $0.pathExtension == "jsonl" },
                scanner: .jsonl(TokenUsageParser.parseClaudeLine)
            ),
            LogSource(
                root: homeDirectory.appending(path: ".codex/sessions"),
                include: { $0.pathExtension == "jsonl" },
                scanner: .jsonl(TokenUsageParser.parseCodexLine)
            ),
            LogSource(
                root: homeDirectory.appending(path: ".gemini"),
                include: { url in
                    url.pathExtension == "jsonl" && url.pathComponents.contains("chats")
                },
                scanner: .jsonl(TokenUsageParser.parseGeminiLine)
            ),
            LogSource(
                root: homeDirectory.appending(path: ".hermes"),
                include: {
                    $0.lastPathComponent == "state.db" &&
                        $0.deletingLastPathComponent().lastPathComponent == ".hermes"
                },
                scanner: .hermesStateDB
            ),
            LogSource(
                root: homeDirectory.appending(path: ".local/share/opencode"),
                include: {
                    $0.lastPathComponent == "opencode.db" &&
                        $0.deletingLastPathComponent().lastPathComponent == "opencode"
                },
                scanner: .openCodeDB
            )
        ]
    }

    private func logFileSnapshots() throws -> [LogFileSnapshot] {
        var snapshots: [LogFileSnapshot] = []
        for source in logSources {
            for file in files(root: source.root, include: source.include) {
                let metadata = try fileMetadata(file)
                snapshots.append(LogFileSnapshot(
                    url: file,
                    path: file.path,
                    size: metadata.size,
                    modifiedAt: metadata.modifiedAt,
                    scanner: source.scanner
                ))
            }
        }
        return snapshots.sorted { $0.path < $1.path }
    }

    private func fileMetadata(_ file: URL) throws -> FileMetadata {
        let values = try file.resourceValues(forKeys: [.fileSizeKey, .contentModificationDateKey])
        return FileMetadata(
            size: Int64(values.fileSize ?? 0),
            modifiedAt: values.contentModificationDate?.timeIntervalSince1970 ?? 0
        )
    }

    private func isSameDisplayWindow(_ lhs: Date, _ rhs: Date) -> Bool {
        var calendar = Calendar.current
        calendar.firstWeekday = 2
        return calendar.isDate(lhs, equalTo: rhs, toGranularity: .hour)
    }
}

private struct LogSource {
    let root: URL
    let include: (URL) -> Bool
    let scanner: SourceScanner
}

private enum SourceScanner {
    case jsonl((String, String) -> UsageRecord?)
    case hermesStateDB
    case openCodeDB

    var supportsIncrementalScan: Bool {
        if case .jsonl = self { return true }
        return false
    }

    var supportsZeroRecordCache: Bool {
        if case .jsonl = self { return true }
        return false
    }
}

private struct ScanProgress {
    let scannedOffset: Int64
    let recordCount: Int
}

private struct FileMetadata {
    let size: Int64
    let modifiedAt: TimeInterval
}

private struct LogFileSnapshot {
    let url: URL
    let path: String
    let size: Int64
    let modifiedAt: TimeInterval
    let scanner: SourceScanner
}

private struct FileSignature: Equatable {
    let entries: [Entry]

    init(snapshots: [LogFileSnapshot]) {
        entries = snapshots.map {
            Entry(path: $0.path, size: $0.size, modifiedAt: $0.modifiedAt)
        }
    }

    struct Entry: Equatable {
        let path: String
        let size: Int64
        let modifiedAt: TimeInterval
    }
}