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
}

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
}

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

    var key: String
    var value: String
}

// MARK: - DatabaseService

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

    private let dbQueue: DatabaseQueue
    private let isoFormatter: ISO8601DateFormatter

    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()
            }
        }

        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: - Conversation Operations

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

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

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

        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: convId,
            name: name,
            messages: savedMessages,
            createdAt: now,
            updatedAt: now
        )
    }

    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
                )
            }

            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
            )

            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
        }
    }
}