sleepy_agent_ios/SleepyAgent/Services/ConversationStorage.swift

import Foundation

// MARK: - Errors

enum ConversationStorageError: LocalizedError {
    case saveFailed(Error)
    case loadFailed(Error)
    case deleteFailed(Error)
    case invalidData
    case conversationNotFound
    
    var errorDescription: String? {
        switch self {
        case .saveFailed(let error):
            return "Failed to save conversation: \(error.localizedDescription)"
        case .loadFailed(let error):
            return "Failed to load conversation: \(error.localizedDescription)"
        case .deleteFailed(let error):
            return "Failed to delete conversation: \(error.localizedDescription)"
        case .invalidData:
            return "Invalid conversation data"
        case .conversationNotFound:
            return "Conversation not found"
        }
    }
}

// MARK: - Data Models

/// Represents a single message in a conversation.
struct ConversationMessage: Codable, Equatable, Identifiable {
    let id: String
    let text: String
    let isUser: Bool
    let isToolCall: Bool
    let timestamp: Date
    
    init(
        id: String = UUID().uuidString,
        text: String,
        isUser: Bool,
        isToolCall: Bool = false,
        timestamp: Date = Date()
    ) {
        self.id = id
        self.text = text
        self.isUser = isUser
        self.isToolCall = isToolCall
        self.timestamp = timestamp
    }
}

/// Metadata about a saved conversation.
struct ConversationInfo: Codable, Equatable, Identifiable {
    let id: String
    let title: String
    let timestamp: Date
    let messageCount: Int
}

/// Full conversation data for storage.
struct SavedConversation: Codable {
    let id: String
    let title: String
    let timestamp: Date
    let messages: [ConversationMessage]
    let version: Int
    
    init(
        id: String,
        title: String,
        timestamp: Date = Date(),
        messages: [ConversationMessage],
        version: Int = 1
    ) {
        self.id = id
        self.title = title
        self.timestamp = timestamp
        self.messages = messages
        self.version = version
    }
}

// MARK: - Storage

/// Manages persistence of conversations using JSON files.
actor ConversationStorage {
    
    // MARK: - Properties
    
    private let fileManager: FileManager
    private let encoder: JSONEncoder
    private let decoder: JSONDecoder
    
    /// Directory where conversations are stored
    let storageDirectory: URL
    
    /// Maximum number of conversations to keep
    var maxConversations: Int = 50
    
    // MARK: - Initialization
    
    init(
        fileManager: FileManager = .default,
        storageDirectory: URL? = nil
    ) {
        self.fileManager = fileManager
        
        // Setup encoder/decoder
        self.encoder = JSONEncoder()
        self.encoder.outputFormatting = [.prettyPrinted, .sortedKeys]
        self.encoder.dateEncodingStrategy = .iso8601
        
        self.decoder = JSONDecoder()
        self.decoder.dateDecodingStrategy = .iso8601
        
        // Setup storage directory
        if let directory = storageDirectory {
            self.storageDirectory = directory
        } else {
            let documentsDirectory = fileManager.urls(for: .documentDirectory, in: .userDomainMask).first!
            self.storageDirectory = documentsDirectory.appendingPathComponent("conversations", isDirectory: true)
        }
        
        // Create directory if needed
        try? fileManager.createDirectory(at: self.storageDirectory, withIntermediateDirectories: true)
    }
    
    // MARK: - Public Methods
    
    /// Creates a new unique conversation ID.
    func createNewConversationId() -> String {
        UUID().uuidString
    }
    
    /// Saves a conversation to storage.
    /// - Parameters:
    ///   - id: Conversation identifier
    ///   - messages: List of messages to save
    /// - Throws: ConversationStorageError if save fails
    func saveConversation(id: String, messages: [ConversationMessage]) throws {
        let title = generateTitle(from: messages)
        let conversation = SavedConversation(
            id: id,
            title: title,
            timestamp: Date(),
            messages: messages
        )
        
        let fileURL = storageDirectory.appendingPathComponent("\(id).json")
        
        do {
            let data = try encoder.encode(conversation)
            try data.write(to: fileURL, options: [.atomic, .completeFileProtection])
            
            // Clean up old conversations if needed
            try cleanupOldConversations()
        } catch {
            throw ConversationStorageError.saveFailed(error)
        }
    }
    
    /// Loads a conversation from storage.
    /// - Parameter id: Conversation identifier
    /// - Returns: Array of messages, or nil if not found
    /// - Throws: ConversationStorageError if load fails
    func loadConversation(id: String) throws -> [ConversationMessage]? {
        let fileURL = storageDirectory.appendingPathComponent("\(id).json")
        
        guard fileManager.fileExists(atPath: fileURL.path) else {
            return nil
        }
        
        do {
            let data = try Data(contentsOf: fileURL)
            let conversation = try decoder.decode(SavedConversation.self, from: data)
            return conversation.messages
        } catch {
            throw ConversationStorageError.loadFailed(error)
        }
    }
    
    /// Deletes a conversation.
    /// - Parameter id: Conversation identifier
    /// - Throws: ConversationStorageError if delete fails
    func deleteConversation(id: String) throws {
        let fileURL = storageDirectory.appendingPathComponent("\(id).json")
        
        guard fileManager.fileExists(atPath: fileURL.path) else {
            throw ConversationStorageError.conversationNotFound
        }
        
        do {
            try fileManager.removeItem(at: fileURL)
        } catch {
            throw ConversationStorageError.deleteFailed(error)
        }
    }
    
    /// Gets metadata for all saved conversations, sorted by most recent.
    /// - Returns: Array of conversation info
    func getAllConversations() -> [ConversationInfo] {
        do {
            let files = try fileManager.contentsOfDirectory(at: storageDirectory, includingPropertiesForKeys: nil)
            let jsonFiles = files.filter { $0.pathExtension == "json" }
            
            return jsonFiles.compactMap { url -> ConversationInfo? in
                guard let conversation = try? loadConversationMetadata(from: url) else { return nil }
                return ConversationInfo(
                    id: conversation.id,
                    title: conversation.title,
                    timestamp: conversation.timestamp,
                    messageCount: conversation.messages.count
                )
            }.sorted { $0.timestamp > $1.timestamp }
        } catch {
            return []
        }
    }
    
    /// Checks if a conversation exists.
    func conversationExists(id: String) -> Bool {
        let fileURL = storageDirectory.appendingPathComponent("\(id).json")
        return fileManager.fileExists(atPath: fileURL.path)
    }
    
    /// Gets the total storage size used by conversations.
    func storageSize() -> Int64 {
        do {
            let files = try fileManager.contentsOfDirectory(at: storageDirectory, includingPropertiesForKeys: [.fileSizeKey])
            var totalSize: Int64 = 0
            for file in files where file.pathExtension == "json" {
                let attributes = try fileManager.attributesOfItem(atPath: file.path)
                totalSize += attributes[.size] as? Int64 ?? 0
            }
            return totalSize
        } catch {
            return 0
        }
    }
    
    /// Exports a conversation to a shareable format.
    func exportConversation(id: String) throws -> Data {
        let fileURL = storageDirectory.appendingPathComponent("\(id).json")
        return try Data(contentsOf: fileURL)
    }
    
    /// Imports a conversation from exported data.
    func importConversation(data: Data) throws -> String {
        let conversation = try decoder.decode(SavedConversation.self, from: data)
        let newId = createNewConversationId()
        
        let importedConversation = SavedConversation(
            id: newId,
            title: conversation.title,
            timestamp: Date(),
            messages: conversation.messages
        )
        
        let fileURL = storageDirectory.appendingPathComponent("\(newId).json")
        let encodedData = try encoder.encode(importedConversation)
        try encodedData.write(to: fileURL, options: [.atomic, .completeFileProtection])
        
        return newId
    }
    
    /// Clears all conversations.
    func clearAllConversations() throws {
        let files = try fileManager.contentsOfDirectory(at: storageDirectory, includingPropertiesForKeys: nil)
        for file in files where file.pathExtension == "json" {
            try fileManager.removeItem(at: file)
        }
    }
    
    // MARK: - Private Methods
    
    private func loadConversationMetadata(from url: URL) throws -> SavedConversation {
        let data = try Data(contentsOf: url)
        return try decoder.decode(SavedConversation.self, from: data)
    }
    
    private func generateTitle(from messages: [ConversationMessage]) -> String {
        let maxLength = 50
        
        // Use first user message as title
        if let firstUserMessage = messages.first(where: { $0.isUser }) {
            let text = firstUserMessage.text.trimmingCharacters(in: .whitespacesAndNewlines)
            if text.count > maxLength {
                let index = text.index(text.startIndex, offsetBy: maxLength)
                return String(text[..<index]) + "..."
            }
            return text
        }
        
        // Fallback to first message or default
        if let firstMessage = messages.first {
            return firstMessage.text.prefix(maxLength).trimmingCharacters(in: .whitespacesAndNewlines)
        }
        
        return "New Chat"
    }
    
    private func cleanupOldConversations() throws {
        let conversations = getAllConversations()
        
        if conversations.count > maxConversations {
            let toDelete = conversations.dropFirst(maxConversations)
            for info in toDelete {
                try? deleteConversation(id: info.id)
            }
        }
    }
}

// MARK: - Conversation Observer

/// Observes conversation storage changes.
actor ConversationObserver {
    private var continuations: [UUID: AsyncStream<[ConversationInfo]>.Continuation] = [:]
    private let storage: ConversationStorage
    private var lastKnownConversations: [ConversationInfo] = []
    private var timer: Timer?
    
    init(storage: ConversationStorage) {
        self.storage = storage
    }
    
    /// Creates a stream that emits when conversations change.
    func observe() -> AsyncStream<[ConversationInfo]> {
        AsyncStream { continuation in
            let id = UUID()
            continuations[id] = continuation
            
            // Initial emit
            Task {
                let conversations = await storage.getAllConversations()
                continuation.yield(conversations)
            }
            
            continuation.onTermination = { [weak self] _ in
                Task {
                    await self?.removeContinuation(id: id)
                }
            }
        }
    }
    
    private func removeContinuation(id: UUID) {
        continuations.removeValue(forKey: id)
    }
    
    /// Manually trigger a check for changes.
    func checkForChanges() async {
        let current = await storage.getAllConversations()
        if current != lastKnownConversations {
            lastKnownConversations = current
            for continuation in continuations.values {
                continuation.yield(current)
            }
        }
    }
}

// MARK: - Preview Helpers

extension ConversationStorage {
    /// Creates a storage instance with sample data for previews.
    static func preview() -> ConversationStorage {
        let storage = ConversationStorage()
        
        // Add sample conversations in background
        Task {
            let sampleMessages = [
                ConversationMessage(text: "Hello, how can you help me?", isUser: true),
                ConversationMessage(text: "I'm an AI assistant. I can help with various tasks like answering questions, searching the web, or having conversations.", isUser: false)
            ]
            try? await storage.saveConversation(
                id: "preview-1",
                messages: sampleMessages
            )
        }
        
        return storage
    }
}