Termind/Frameworks/TermindApplication/FeatureModels/MemoryModel.swift

import AgentKit
import Foundation
import TermindDomain

@Observable @MainActor
public final class MemoryModel {
    public private(set) var isEnabled: Bool
    public private(set) var isBusy = false
    public private(set) var isIndexing = false
    public private(set) var state = AgentMemoryState()
    public private(set) var error: String?
    public var reviewEnabled: Bool {
        didSet {
            defaults?.set(reviewEnabled, forKey: Self.reviewKey)
            if !reviewEnabled {
                for task in reviews.values { task.cancel() }
                pendingReviews.removeAll()
            }
        }
    }
    public var reviewModel: AIModelRef? {
        didSet { defaults?.set(try? JSONEncoder().encode(reviewModel), forKey: Self.modelKey) }
    }
    private static let enabledKey = "hermesMemoryEnabled"
    private static let destroyingKey = "hermesMemoryDestroyPending"
    private static let reviewKey = "hermesMemoryReviewEnabled"
    private static let modelKey = "hermesMemoryReviewModel"
    private let defaults: UserDefaults?
    private let storage: (any MemoryPersisting)?
    private var generation = UUID()
    private var contexts: [UUID: AgentMemoryContext] = [:]
    private var history: (any AgentRunPersisting)?
    private var indexTask: Task<Void, Never>?
    private var indexDirty = false
    private var reviews: [UUID: Task<Void, Never>] = [:]
    private var pendingReviews: [UUID: (@Sendable () async -> Void)] = [:]

    public init(storage: (any MemoryPersisting)? = nil, defaults: UserDefaults? = nil) {
        self.storage = storage
        self.defaults = defaults
        let pending = defaults?.bool(forKey: Self.destroyingKey) == true
        isEnabled = !pending && storage != nil && defaults?.bool(forKey: Self.enabledKey) == true
        reviewEnabled = defaults?.object(forKey: Self.reviewKey) as? Bool ?? true
        reviewModel = defaults?.data(forKey: Self.modelKey).flatMap {
            try? JSONDecoder().decode(AIModelRef.self, from: $0)
        }
        if pending { Task { await destroy() } }
    }
    public func attachHistory(_ history: any AgentRunPersisting) {
        self.history = history
        historyChanged()
    }
    private func invalidate() {
        generation = UUID()
        for context in contexts.values { context.invalidate() }
        contexts.removeAll()
        indexTask?.cancel()
        for task in reviews.values { task.cancel() }
        pendingReviews.removeAll()
    }
    public func setEnabled(_ enabled: Bool) async {
        guard !isBusy, let storage, defaults?.bool(forKey: Self.destroyingKey) != true else { return }
        isBusy = true
        defer { isBusy = false }
        invalidate()
        isEnabled = enabled
        defaults?.set(enabled, forKey: Self.enabledKey)
        error = nil
        do {
            if enabled {
                state = try await storage.memoryState()
                historyChanged()
            } else {
                try await storage.close()
            }
        } catch { self.error = error.localizedDescription }
    }
    public func refresh() async {
        guard isEnabled, let storage else { return }
        let token = generation
        do {
            let value = try await storage.memoryState()
            if token == generation {
                state = value
                error = nil
            }
        } catch { if token == generation { self.error = error.localizedDescription } }
    }
    public func edit(_ target: AgentMemoryTarget, text: String) async {
        guard isEnabled, let storage else { return }
        let token = generation
        do {
            let updated = try await storage.replaceTarget(target, text: text)
            guard token == generation, isEnabled else { return }
            state = updated
            error = nil
        } catch { if token == generation { self.error = error.localizedDescription } }
    }
    public func session(_ id: UUID) async -> (context: AgentMemoryContext, access: any AgentMemoryAccessing)? {
        guard isEnabled, let storage else { return nil }
        let token = generation
        do {
            let context: AgentMemoryContext
            if let existing = contexts[id] {
                context = existing
            } else {
                let snapshot = try await storage.snapshot(for: id)
                guard token == generation, isEnabled else { return nil }
                context = contexts[id] ?? AgentMemoryContext(snapshot: snapshot)
                contexts[id] = context
            }
            return (context, MemoryAccess(model: self, token: token))
        } catch {
            if token == generation { self.error = error.localizedDescription }
            return nil
        }
    }
    private func require(_ token: UUID) throws -> any MemoryPersisting {
        try Task.checkCancellation()
        guard isEnabled, token == generation, !isBusy, let storage else { throw CancellationError() }
        return storage
    }
    fileprivate func read(token: UUID) async throws -> AgentMemoryState {
        let result = try await require(token).memoryState()
        _ = try require(token)
        return result
    }
    @discardableResult
    fileprivate func apply(_ operations: [AgentMemoryOperation], token: UUID) async throws -> AgentMemoryState {
        let result = try await require(token).applyMemory(operations)
        _ = try require(token)
        state = result
        return result
    }
    fileprivate func search(_ request: AgentSessionSearchRequest, token: UUID) async throws -> AgentJSONValue {
        _ = try require(token)
        historyChanged()
        await indexTask?.value
        let result = try await require(token).searchSessions(request)
        _ = try require(token)
        return result
    }
    public func rebuildHistoryIndex() async {
        guard isEnabled else { return }
        let token = generation
        historyChanged()
        await indexTask?.value
        do {
            try await require(token).rebuildHistoryIndex()
            if token == generation { error = nil }
        } catch { if token == generation { self.error = error.localizedDescription } }
    }

    public func historyChanged() {
        guard isEnabled, let storage, let history else { return }
        indexDirty = true
        guard indexTask == nil else { return }
        let token = generation
        isIndexing = true
        indexTask = Task { [weak self] in
            guard let self else { return }
            defer {
                indexTask = nil
                isIndexing = false
                if indexDirty, isEnabled { historyChanged() }
            }
            do {
                while indexDirty, token == generation, isEnabled, !Task.isCancelled {
                    indexDirty = false
                    let conversations = try await history.loadConversations()
                    _ = try require(token)
                    try await storage.replaceHistory(conversations)
                }
            } catch is CancellationError {
            } catch {
                if token == generation {
                    self.error = error.localizedDescription
                    indexDirty = false
                }
            }
        }
    }
    public func review(
        conversationID: UUID, messages: [AgentTranscriptMessage], model: any AgentModelCompleting,
        wait: Bool = false
    ) async {
        guard isEnabled, reviewEnabled, let storage, let last = messages.last?.id else { return }
        let token = generation
        if let running = reviews[conversationID] {
            if wait {
                await running.value
            } else {
                pendingReviews[conversationID] = { [weak self] in
                    await self?.review(conversationID: conversationID, messages: messages, model: model)
                }
                return
            }
        }
        guard token == generation, isEnabled else { return }
        let task = Task { [weak self] in
            guard let self else { return }
            defer {
                reviews.removeValue(forKey: conversationID)
                if let next = pendingReviews.removeValue(forKey: conversationID) { Task { await next() } }
            }
            do {
                guard try await storage.lastReviewedMessage(conversationID) != last else { return }
                let current = try await read(token: token)
                let operations = try await withThrowingTaskGroup(of: [AgentMemoryOperation].self) { group in
                    group.addTask {
                        try await AgentMemoryLearningService().review(messages: messages, state: current, model: model)
                    }
                    group.addTask {
                        try await Task.sleep(for: .seconds(60))
                        throw CancellationError()
                    }
                    defer { group.cancelAll() }
                    return try await group.next() ?? []
                }
                _ = try require(token)
                if let updated = try await storage.applyReview(
                    operations, expected: current, messageID: last, conversationID: conversationID)
                {
                    _ = try require(token)
                    state = updated
                }
            } catch is CancellationError {
            } catch { if token == generation { self.error = error.localizedDescription } }
        }
        reviews[conversationID] = task
        if wait { await task.value }
    }
    public func destroy() async {
        guard !isBusy, let storage else { return }
        isBusy = true
        isEnabled = false
        defaults?.set(false, forKey: Self.enabledKey)
        defaults?.set(true, forKey: Self.destroyingKey)
        invalidate()
        indexDirty = false
        let pending = Array(reviews.values)
        await indexTask?.value
        for task in pending { await task.value }
        do {
            try await storage.destroy()
            state = .init()
            error = nil
            defaults?.removeObject(forKey: Self.destroyingKey)
        } catch { self.error = error.localizedDescription }
        isBusy = false
    }
    public func suspend() {
        indexTask?.cancel()
        for task in reviews.values { task.cancel() }
        pendingReviews.removeAll()
    }
}

private struct MemoryAccess: AgentMemoryAccessing {
    weak var model: MemoryModel?
    let token: UUID
    func memoryState() async throws -> AgentMemoryState {
        guard let model else { throw CancellationError() }
        return try await model.read(token: token)
    }
    func applyMemory(_ operations: [AgentMemoryOperation]) async throws -> AgentMemoryState {
        guard let model else { throw CancellationError() }
        return try await model.apply(operations, token: token)
    }
    func searchSessions(_ request: AgentSessionSearchRequest) async throws -> AgentJSONValue {
        guard let model else { throw CancellationError() }
        return try await model.search(request, token: token)
    }
}