gramora/gramora/Services/ParaphrasingService.swift

import Foundation

protocol ParaphrasingServicing {
    func paraphrase(_ text: String) async throws -> String
}

struct ParaphrasingService: ParaphrasingServicing {
    private static let chunkMaxWords = 280
    private static let maxConcurrentChunks = 4
    private static let systemPrompt = """
    You are a professional writing assistant. Paraphrase the user's text while preserving \
    the original meaning and tone. Return only the paraphrased text with no preamble or explanation.
    """

    private let openAI: OpenAIChatServicing

    init(openAI: OpenAIChatServicing = OpenAIChatService.shared) {
        self.openAI = openAI
    }

    func paraphrase(_ text: String) async throws -> String {
        let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !trimmed.isEmpty else {
            throw OpenAIServiceError.apiError("Enter some text to paraphrase.")
        }

        let chunks = trimmed.chunked(maxWords: Self.chunkMaxWords)
        if chunks.count == 1 {
            return try await openAI.completeText(
                systemPrompt: Self.systemPrompt,
                userPrompt: trimmed,
                maxTokens: 1_000
            )
        }

        var paraphrasedChunks: [String] = []

        for batchStart in stride(from: 0, to: chunks.count, by: Self.maxConcurrentChunks) {
            let batchEnd = min(batchStart + Self.maxConcurrentChunks, chunks.count)
            let batch = Array(chunks[batchStart..<batchEnd])

            let batchResults = try await withThrowingTaskGroup(of: (Int, String).self) { group in
                for (offset, chunk) in batch.enumerated() {
                    let originalIndex = batchStart + offset
                    group.addTask {
                        let paraphrased = try await self.openAI.completeText(
                            systemPrompt: Self.systemPrompt,
                            userPrompt: chunk,
                            maxTokens: 500
                        )
                        return (originalIndex, paraphrased)
                    }
                }

                var indexedResults: [(Int, String)] = []
                for try await item in group {
                    indexedResults.append(item)
                }
                return indexedResults.sorted { $0.0 < $1.0 }
            }

            paraphrasedChunks.append(contentsOf: batchResults.map(\.1))
        }

        return paraphrasedChunks.joined(separator: "\n\n")
    }
}