oai-swift/oAI/Services/ConversationMergeService.swift

//
//  ConversationMergeService.swift
//  oAI
//
//  Combine multiple saved conversations into one (simple concatenation or AI-assisted merge)
//
// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright (C) 2026 Rune Olsen
//
// This file is part of oAI.
//
// oAI is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// oAI is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
// or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General
// Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with oAI. If not, see <https://www.gnu.org/licenses/>.


import Foundation
import os

enum CombineMode: String, Sendable {
    case simple
    case ai
}

enum MergeError: LocalizedError {
    case tooFewConversations
    case noDefaultModel
    case noAPIKey
    case invalidAIResponse(String)

    var errorDescription: String? {
        switch self {
        case .tooFewConversations:
            return "Select at least two conversations to combine."
        case .noDefaultModel:
            return "No default model is configured. Set one in Settings → General → Default Model."
        case .noAPIKey:
            return "No API key configured for the default provider. Add one in Settings."
        case .invalidAIResponse(let snippet):
            return "The model's response could not be parsed into a conversation: \(snippet)"
        }
    }
}

enum ConversationMergeService {

    static func merge(
        conversationIds: [UUID],
        name: String,
        mode: CombineMode,
        deleteOriginals: Bool
    ) async throws -> Conversation {
        guard conversationIds.count >= 2 else {
            throw MergeError.tooFewConversations
        }

        let sources: [(Conversation, [Message])] = try conversationIds.compactMap { id in
            try DatabaseService.shared.loadConversation(id: id)
        }

        let mergedMessages: [Message]
        switch mode {
        case .simple:
            mergedMessages = simpleMerge(sources)
        case .ai:
            mergedMessages = try await aiMerge(sources)
        }

        let newConversation = try DatabaseService.shared.saveConversation(name: name, messages: mergedMessages)

        if deleteOriginals {
            for id in conversationIds {
                _ = try? DatabaseService.shared.deleteConversation(id: id)
            }
        }

        Log.db.info("Combined \(conversationIds.count) conversations into '\(name)' (mode: \(mode.rawValue), deleteOriginals: \(deleteOriginals))")

        return newConversation
    }

    private static func simpleMerge(_ sources: [(Conversation, [Message])]) -> [Message] {
        sources.flatMap { $0.1 }.sorted { $0.timestamp < $1.timestamp }
    }

    private struct MergedTurn: Codable {
        let role: String
        let content: String
    }

    private static func aiMerge(_ sources: [(Conversation, [Message])]) async throws -> [Message] {
        let settings = SettingsService.shared
        guard let modelId = settings.defaultModel, !modelId.isEmpty else {
            throw MergeError.noDefaultModel
        }
        guard let provider = ProviderRegistry.shared.getProvider(for: settings.defaultProvider) else {
            throw MergeError.noAPIKey
        }

        let transcript = sources.map { conversation, messages -> String in
            let body = messages.map { msg -> String in
                let label = msg.role == .user ? "**User:**" : "**Assistant:**"
                return "\(label) \(msg.content)"
            }.joined(separator: "\n\n")
            return "### Conversation: \(conversation.name)\n\n\(body)"
        }.joined(separator: "\n\n---\n\n")

        let mergePrompt = """
        Merge the following saved conversation transcripts into a single, coherent conversation. \
        Remove redundant or duplicate exchanges, keep the most informative answer when sources overlap, \
        preserve important details from each source, and do not invent facts that were not in the originals.

        Respond with ONLY a JSON array of message objects in logical order, each in the form \
        {"role": "user" or "assistant", "content": "..."}. Do not include any text outside the JSON array.

        \(transcript)
        """

        let request = ChatRequest(
            messages: [Message(role: .user, content: mergePrompt)],
            model: modelId,
            stream: false,
            maxTokens: 4000,
            temperature: 0.3,
            topP: nil,
            systemPrompt: "You are a helpful assistant that merges chat conversation transcripts into one clean, coherent conversation.",
            tools: nil,
            onlineMode: false,
            imageGeneration: false
        )

        let response: ChatResponse
        do {
            response = try await provider.chat(request: request)
        } catch {
            Log.api.error("Conversation merge AI call failed: \(error.localizedDescription)")
            throw error
        }

        let turns = try parseTurns(from: response.content)

        let base = Date()
        return turns.enumerated().map { index, turn in
            Message(
                role: turn.role == "user" ? .user : .assistant,
                content: turn.content,
                timestamp: base.addingTimeInterval(TimeInterval(index)),
                modelId: modelId
            )
        }
    }

    private static func parseTurns(from raw: String) throws -> [MergedTurn] {
        var text = raw.trimmingCharacters(in: .whitespacesAndNewlines)
        if text.hasPrefix("```") {
            text = text.components(separatedBy: "\n").dropFirst().joined(separator: "\n")
            if text.hasSuffix("```") {
                text = String(text.dropLast(3))
            }
            text = text.trimmingCharacters(in: .whitespacesAndNewlines)
        }
        guard let data = text.data(using: .utf8),
              let turns = try? JSONDecoder().decode([MergedTurn].self, from: data),
              !turns.isEmpty else {
            throw MergeError.invalidAIResponse(String(raw.prefix(200)))
        }
        return turns
    }
}