oai-swift/oAI/Services/DatabaseService.swift

//
//  DatabaseService.swift
//  oAI
//
//  SQLite persistence layer for conversations using GRDB
//

import Foundation
import GRDB
import os

// MARK: - Database Record Types

struct ConversationRecord: Codable, FetchableRecord, PersistableRecord, Sendable {
    static let databaseTableName = "conversations"

    var id: String
    var name: String
    var createdAt: String
    var updatedAt: String
    var primaryModel: String?
}

struct MessageRecord: Codable, FetchableRecord, PersistableRecord, Sendable {
    static let databaseTableName = "messages"

    var id: String
    var conversationId: String
    var role: String
    var content: String
    var tokens: Int?
    var cost: Double?
    var timestamp: String
    var sortOrder: Int
    var modelId: String?
}

struct SettingRecord: Codable, FetchableRecord, PersistableRecord, Sendable {
    static let databaseTableName = "settings"

    var key: String
    var value: String
}

struct HistoryRecord: Codable, FetchableRecord, PersistableRecord, Sendable {
    static let databaseTableName = "command_history"

    var id: String
    var input: String
    var timestamp: String
}

struct EmailLogRecord: Codable, FetchableRecord, PersistableRecord, Sendable {
    static let databaseTableName = "email_logs"

    var id: String
    var timestamp: String
    var sender: String
    var subject: String
    var emailContent: String
    var aiResponse: String?
    var status: String // "success" or "error"
    var errorMessage: String?
    var tokens: Int?
    var cost: Double?
    var responseTime: Double?
    var modelId: String?
}

// MARK: - DatabaseService

final class DatabaseService: Sendable {
    nonisolated static let shared = DatabaseService()

    private let dbQueue: DatabaseQueue
    private let isoFormatter: ISO8601DateFormatter

    // Command history limit - keep most recent 5000 entries
    private static let maxHistoryEntries = 5000

    nonisolated private init() {
        isoFormatter = ISO8601DateFormatter()
        isoFormatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]

        let fileManager = FileManager.default
        let appSupport = fileManager.urls(for: .applicationSupportDirectory, in: .userDomainMask).first!
        let dbDirectory = appSupport.appendingPathComponent("oAI", isDirectory: true)

        try! fileManager.createDirectory(at: dbDirectory, withIntermediateDirectories: true)

        let dbPath = dbDirectory.appendingPathComponent("oai_conversations.db").path
        Log.db.info("Opening database at \(dbPath)")
        dbQueue = try! DatabaseQueue(path: dbPath)

        try! migrator.migrate(dbQueue)
    }

    private var migrator: DatabaseMigrator {
        var migrator = DatabaseMigrator()

        migrator.registerMigration("v1") { db in
            try db.create(table: "conversations") { t in
                t.primaryKey("id", .text)
                t.column("name", .text).notNull()
                t.column("createdAt", .text).notNull()
                t.column("updatedAt", .text).notNull()
            }

            try db.create(table: "messages") { t in
                t.primaryKey("id", .text)
                t.column("conversationId", .text).notNull()
                    .references("conversations", onDelete: .cascade)
                t.column("role", .text).notNull()
                t.column("content", .text).notNull()
                t.column("tokens", .integer)
                t.column("cost", .double)
                t.column("timestamp", .text).notNull()
                t.column("sortOrder", .integer).notNull()
            }

            try db.create(
                index: "messages_on_conversationId",
                on: "messages",
                columns: ["conversationId"]
            )
        }

        migrator.registerMigration("v2") { db in
            try db.create(table: "settings") { t in
                t.primaryKey("key", .text)
                t.column("value", .text).notNull()
            }
        }

        migrator.registerMigration("v3") { db in
            try db.create(table: "command_history") { t in
                t.primaryKey("id", .text)
                t.column("input", .text).notNull()
                t.column("timestamp", .text).notNull()
            }

            try db.create(
                index: "command_history_on_timestamp",
                on: "command_history",
                columns: ["timestamp"]
            )
        }

        migrator.registerMigration("v4") { db in
            // Add modelId to messages table (nullable for existing messages)
            try db.alter(table: "messages") { t in
                t.add(column: "modelId", .text)
            }

            // Add primaryModel to conversations table (nullable)
            try db.alter(table: "conversations") { t in
                t.add(column: "primaryModel", .text)
            }
        }

        migrator.registerMigration("v5") { db in
            // Email handler logs
            try db.create(table: "email_logs") { t in
                t.primaryKey("id", .text)
                t.column("timestamp", .text).notNull()
                t.column("sender", .text).notNull()
                t.column("subject", .text).notNull()
                t.column("emailContent", .text).notNull()
                t.column("aiResponse", .text)
                t.column("status", .text).notNull() // "success" or "error"
                t.column("errorMessage", .text)
                t.column("tokens", .integer)
                t.column("cost", .double)
                t.column("responseTime", .double)
                t.column("modelId", .text)
            }

            try db.create(
                index: "email_logs_on_timestamp",
                on: "email_logs",
                columns: ["timestamp"]
            )
        }

        return migrator
    }

    // MARK: - Settings Operations

    nonisolated func loadAllSettings() throws -> [String: String] {
        try dbQueue.read { db in
            let records = try SettingRecord.fetchAll(db)
            return Dictionary(uniqueKeysWithValues: records.map { ($0.key, $0.value) })
        }
    }

    nonisolated func setSetting(key: String, value: String) {
        try? dbQueue.write { db in
            let record = SettingRecord(key: key, value: value)
            try record.save(db)
        }
    }

    nonisolated func deleteSetting(key: String) {
        try? dbQueue.write { db in
            _ = try SettingRecord.deleteOne(db, key: key)
        }
    }

    // MARK: - Encrypted Settings Operations

    /// Store an encrypted setting (for sensitive data like API keys)
    nonisolated func setEncryptedSetting(key: String, value: String) throws {
        let encryptedValue = try EncryptionService.shared.encrypt(value)
        try dbQueue.write { db in
            let record = SettingRecord(key: "encrypted_\(key)", value: encryptedValue)
            try record.save(db)
        }
    }

    /// Retrieve and decrypt an encrypted setting
    nonisolated func getEncryptedSetting(key: String) throws -> String? {
        let encryptedValue = try dbQueue.read { db in
            try SettingRecord.fetchOne(db, key: "encrypted_\(key)")?.value
        }

        guard let encryptedValue = encryptedValue else {
            return nil
        }

        return try EncryptionService.shared.decrypt(encryptedValue)
    }

    /// Delete an encrypted setting
    nonisolated func deleteEncryptedSetting(key: String) {
        try? dbQueue.write { db in
            _ = try SettingRecord.deleteOne(db, key: "encrypted_\(key)")
        }
    }

    // MARK: - Conversation Operations

    nonisolated func saveConversation(name: String, messages: [Message]) throws -> Conversation {
        return try saveConversation(id: UUID(), name: name, messages: messages, primaryModel: nil)
    }

    nonisolated func saveConversation(id: UUID, name: String, messages: [Message], primaryModel: String?) throws -> Conversation {
        Log.db.info("Saving conversation '\(name)' with \(messages.count) messages (primaryModel: \(primaryModel ?? "none"))")
        let now = Date()
        let nowString = isoFormatter.string(from: now)

        let convRecord = ConversationRecord(
            id: id.uuidString,
            name: name,
            createdAt: nowString,
            updatedAt: nowString,
            primaryModel: primaryModel
        )

        let messageRecords = messages.enumerated().compactMap { index, msg -> MessageRecord? in
            guard msg.role != .system else { return nil }
            return MessageRecord(
                id: UUID().uuidString,
                conversationId: id.uuidString,
                role: msg.role.rawValue,
                content: msg.content,
                tokens: msg.tokens,
                cost: msg.cost,
                timestamp: isoFormatter.string(from: msg.timestamp),
                sortOrder: index,
                modelId: msg.modelId
            )
        }

        try dbQueue.write { db in
            try convRecord.insert(db)
            for record in messageRecords {
                try record.insert(db)
            }
        }

        let savedMessages = messages.filter { $0.role != .system }
        return Conversation(
            id: id,
            name: name,
            messages: savedMessages,
            createdAt: now,
            updatedAt: now,
            primaryModel: primaryModel
        )
    }

    nonisolated func loadConversation(id: UUID) throws -> (Conversation, [Message])? {
        try dbQueue.read { db in
            guard let convRecord = try ConversationRecord.fetchOne(db, key: id.uuidString) else {
                return nil
            }

            let messageRecords = try MessageRecord
                .filter(Column("conversationId") == id.uuidString)
                .order(Column("sortOrder"))
                .fetchAll(db)

            let messages = messageRecords.compactMap { record -> Message? in
                guard let msgId = UUID(uuidString: record.id),
                      let role = MessageRole(rawValue: record.role),
                      let timestamp = self.isoFormatter.date(from: record.timestamp)
                else { return nil }

                return Message(
                    id: msgId,
                    role: role,
                    content: record.content,
                    tokens: record.tokens,
                    cost: record.cost,
                    timestamp: timestamp,
                    modelId: record.modelId
                )
            }

            guard let convId = UUID(uuidString: convRecord.id),
                  let createdAt = self.isoFormatter.date(from: convRecord.createdAt),
                  let updatedAt = self.isoFormatter.date(from: convRecord.updatedAt)
            else { return nil }

            let conversation = Conversation(
                id: convId,
                name: convRecord.name,
                messages: messages,
                createdAt: createdAt,
                updatedAt: updatedAt,
                primaryModel: convRecord.primaryModel
            )

            return (conversation, messages)
        }
    }

    nonisolated func listConversations() throws -> [Conversation] {
        try dbQueue.read { db in
            let records = try ConversationRecord
                .order(Column("updatedAt").desc)
                .fetchAll(db)

            return records.compactMap { record -> Conversation? in
                guard let id = UUID(uuidString: record.id),
                      let createdAt = self.isoFormatter.date(from: record.createdAt),
                      let updatedAt = self.isoFormatter.date(from: record.updatedAt)
                else { return nil }

                // Fetch message count without loading all messages
                let messageCount = (try? MessageRecord
                    .filter(Column("conversationId") == record.id)
                    .fetchCount(db)) ?? 0

                // Get last message date
                let lastMsg = try? MessageRecord
                    .filter(Column("conversationId") == record.id)
                    .order(Column("sortOrder").desc)
                    .fetchOne(db)

                let lastDate = lastMsg.flatMap { self.isoFormatter.date(from: $0.timestamp) } ?? updatedAt

                // Create conversation with empty messages array but correct metadata
                var conv = Conversation(
                    id: id,
                    name: record.name,
                    messages: Array(repeating: Message(role: .user, content: ""), count: messageCount),
                    createdAt: createdAt,
                    updatedAt: lastDate
                )
                // We store placeholder messages just for the count; lastMessageDate uses updatedAt
                conv.updatedAt = lastDate
                return conv
            }
        }
    }

    nonisolated func deleteConversation(id: UUID) throws -> Bool {
        Log.db.info("Deleting conversation \(id.uuidString)")
        return try dbQueue.write { db in
            try MessageRecord.filter(Column("conversationId") == id.uuidString).deleteAll(db)
            return try ConversationRecord.deleteOne(db, key: id.uuidString)
        }
    }

    nonisolated func deleteConversation(name: String) throws -> Bool {
        try dbQueue.write { db in
            guard let record = try ConversationRecord
                .filter(Column("name") == name)
                .fetchOne(db)
            else { return false }

            try MessageRecord.filter(Column("conversationId") == record.id).deleteAll(db)
            return try ConversationRecord.deleteOne(db, key: record.id)
        }
    }

    nonisolated func updateConversation(id: UUID, name: String?, messages: [Message]?) throws -> Bool {
        try dbQueue.write { db in
            guard var convRecord = try ConversationRecord.fetchOne(db, key: id.uuidString) else {
                return false
            }

            if let name = name {
                convRecord.name = name
            }

            convRecord.updatedAt = self.isoFormatter.string(from: Date())
            try convRecord.update(db)

            if let messages = messages {
                try MessageRecord.filter(Column("conversationId") == id.uuidString).deleteAll(db)

                let messageRecords = messages.enumerated().compactMap { index, msg -> MessageRecord? in
                    guard msg.role != .system else { return nil }
                    return MessageRecord(
                        id: UUID().uuidString,
                        conversationId: id.uuidString,
                        role: msg.role.rawValue,
                        content: msg.content,
                        tokens: msg.tokens,
                        cost: msg.cost,
                        timestamp: self.isoFormatter.string(from: msg.timestamp),
                        sortOrder: index
                    )
                }

                for record in messageRecords {
                    try record.insert(db)
                }
            }

            return true
        }
    }

    // MARK: - Command History Operations

    nonisolated func saveCommandHistory(input: String) {
        let now = Date()
        let record = HistoryRecord(
            id: UUID().uuidString,
            input: input,
            timestamp: isoFormatter.string(from: now)
        )

        try? dbQueue.write { db in
            try record.insert(db)

            // Clean up old entries if we exceed the limit
            let count = try HistoryRecord.fetchCount(db)
            if count > Self.maxHistoryEntries {
                // Delete oldest entries to get back to limit
                let excess = count - Self.maxHistoryEntries
                try db.execute(
                    sql: """
                    DELETE FROM command_history
                    WHERE id IN (
                        SELECT id FROM command_history
                        ORDER BY timestamp ASC
                        LIMIT ?
                    )
                    """,
                    arguments: [excess]
                )
            }
        }
    }

    nonisolated func loadCommandHistory() throws -> [(input: String, timestamp: Date)] {
        try dbQueue.read { db in
            let records = try HistoryRecord
                .order(Column("timestamp").desc)
                .fetchAll(db)

            return records.compactMap { record in
                guard let date = isoFormatter.date(from: record.timestamp) else {
                    return nil
                }
                return (input: record.input, timestamp: date)
            }
        }
    }

    nonisolated func searchCommandHistory(query: String) throws -> [(input: String, timestamp: Date)] {
        try dbQueue.read { db in
            let records = try HistoryRecord
                .filter(Column("input").like("%\(query)%"))
                .order(Column("timestamp").desc)
                .fetchAll(db)

            return records.compactMap { record in
                guard let date = isoFormatter.date(from: record.timestamp) else {
                    return nil
                }
                return (input: record.input, timestamp: date)
            }
        }
    }

    // MARK: - Email Log Operations

    nonisolated func saveEmailLog(_ log: EmailLog) {
        let record = EmailLogRecord(
            id: log.id.uuidString,
            timestamp: isoFormatter.string(from: log.timestamp),
            sender: log.sender,
            subject: log.subject,
            emailContent: log.emailContent,
            aiResponse: log.aiResponse,
            status: log.status.rawValue,
            errorMessage: log.errorMessage,
            tokens: log.tokens,
            cost: log.cost,
            responseTime: log.responseTime,
            modelId: log.modelId
        )

        try? dbQueue.write { db in
            try record.insert(db)
        }
    }

    nonisolated func loadEmailLogs(limit: Int = 100) throws -> [EmailLog] {
        try dbQueue.read { db in
            let records = try EmailLogRecord
                .order(Column("timestamp").desc)
                .limit(limit)
                .fetchAll(db)

            return records.compactMap { record in
                guard let timestamp = isoFormatter.date(from: record.timestamp),
                      let status = EmailLogStatus(rawValue: record.status),
                      let id = UUID(uuidString: record.id) else {
                    return nil
                }

                return EmailLog(
                    id: id,
                    timestamp: timestamp,
                    sender: record.sender,
                    subject: record.subject,
                    emailContent: record.emailContent,
                    aiResponse: record.aiResponse,
                    status: status,
                    errorMessage: record.errorMessage,
                    tokens: record.tokens,
                    cost: record.cost,
                    responseTime: record.responseTime,
                    modelId: record.modelId
                )
            }
        }
    }

    nonisolated func deleteEmailLog(id: UUID) {
        try? dbQueue.write { db in
            try db.execute(
                sql: "DELETE FROM email_logs WHERE id = ?",
                arguments: [id.uuidString]
            )
        }
    }

    nonisolated func clearEmailLogs() {
        try? dbQueue.write { db in
            try db.execute(sql: "DELETE FROM email_logs")
        }
    }

    nonisolated func getEmailLogCount() throws -> Int {
        try dbQueue.read { db in
            try EmailLogRecord.fetchCount(db)
        }
    }
}