//
// Created by Banghua Zhao on 09/07/2025
// Copyright Apps Bay Limited. All rights reserved.
//

import Foundation
import NaturalLanguage

// MARK: - RAG Demo Result

struct RAGDemoResult {
    var slices: [String]?
    var indexed: [(chunk: String, embedding: [Double])]?
    var retrieved: [String]?
    var reranked: [String]?
    var generated: (prompt: String, response: String)?
    var error: Error?
}

// MARK: - RAG Service

class RAGService {
    // MARK: - Dependencies

    private let geminiService: GeminiService

    // MARK: - Local Embedding Models

    private let nlEmbedding = NLEmbedding.sentenceEmbedding(for: .english)!

    init(geminiService: GeminiService = GeminiService()) {
        // Initialize Gemini service
        self.geminiService = geminiService

        // Initialize Natural Language embedding
        print("RAGService initialized with local embedding models")
        print("Embedding dimension: \(nlEmbedding.dimension)")
    }

    // MARK: - Document Processing

    func loadDocument(_ documentURL: URL) throws -> String {
        try String(contentsOf: documentURL)
    }

    func slice(text: String) -> [String] {
        text.components(separatedBy: "\n\n")
            .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) }
            .filter { !$0.isEmpty }
    }

    // MARK: - Main Embedding Function

    func embed(text: String) -> [Double] {
        if let vector = nlEmbedding.vector(for: text) {
            // Normalize the embedding to mimic text2vec's normalized output
            let norm = sqrt(vector.map { $0 * $0 }.reduce(0, +))
            return norm > 0 ? vector.map { $0 / norm } : Array(vector)
        }
        // Fallback to zero vector
        return Array(repeating: 0.0, count: nlEmbedding.dimension)
    }

    func embed(texts: [String]) -> [[Double]] {
        return texts.map { embed(text: $0) }
    }

    func indexWithEmbeddings(slices: [String]) -> [(chunk: String, embedding: [Double])] {
        let embeddings = embed(texts: slices)
        guard embeddings.count == slices.count else {
            return []
        }
        return Array(zip(slices, embeddings))
    }

    func cosineSimilarity(_ a: [Double], _ b: [Double]) -> Double {
        guard a.count == b.count else { return -1 }
        let dot = zip(a, b).map(*).reduce(0, +)
        let normA = sqrt(a.map { $0 * $0 }.reduce(0, +))
        let normB = sqrt(b.map { $0 * $0 }.reduce(0, +))
        return dot / (normA * normB)
    }

    func retrieve(index: [(chunk: String, embedding: [Double])], query: String, topK: Int = 10) -> [String] {
        let queryEmbedding = embed(text: query)
        let scored = index.map { ($0.chunk, self.cosineSimilarity($0.embedding, queryEmbedding)) }
        let sorted = scored.sorted { $0.1 > $1.1 }
        return Array(sorted.prefix(topK).map { $0.0 })
    }

    /// Mimic cross-encoder by combining semantic similarity with contextual features
    func rerank(chunks: [String], query: String, topK: Int = 3) -> [String] {
        let queryEmbedding = embed(text: query)
        let queryWords = Set(query.lowercased().components(separatedBy: .whitespacesAndNewlines).filter { !$0.isEmpty })

        let scored = chunks.map { chunk -> (String, Double) in
            // Semantic similarity score
            let chunkEmbedding = embed(text: chunk)
            let semanticScore = cosineSimilarity(chunkEmbedding, queryEmbedding)

            // Word overlap score
            let chunkWords = Set(chunk.lowercased().components(separatedBy: .whitespacesAndNewlines).filter { !$0.isEmpty })
            let overlap = Double(queryWords.intersection(chunkWords).count)
            let overlapScore = overlap / Double(max(queryWords.count, 1))

            // Length normalization (penalize overly long or short chunks)
            let wordCount = chunkWords.count
            let lengthPenalty = min(1.0, max(0.5, Double(wordCount) / 50.0)) // Normalize around 50 words

            // Proximity bonus: check if query terms appear close together in chunk
            let proximityScore = computeProximityScore(chunk: chunk.lowercased(), queryWords: queryWords)

            // Combine scores (weights tuned to prioritize semantic similarity)
            let finalScore = 0.6 * semanticScore + 0.2 * overlapScore + 0.1 * lengthPenalty + 0.1 * proximityScore
            return (chunk, finalScore)
        }

        return scored.sorted { $0.1 > $1.1 }.map { $0.0 }.prefix(topK).map { String($0) }
    }

    // Helper function to compute proximity score
    private func computeProximityScore(chunk: String, queryWords: Set<String>) -> Double {
        let words = chunk.components(separatedBy: .whitespacesAndNewlines).filter { !$0.isEmpty }
        var minDistance = Double.infinity
        let queryWordArray = Array(queryWords)

        // Check pairwise distances between query words in the chunk
        for i in 0 ..< queryWordArray.count {
            for j in (i + 1) ..< queryWordArray.count {
                if let iPos = words.firstIndex(of: queryWordArray[i].lowercased()),
                   let jPos = words.firstIndex(of: queryWordArray[j].lowercased()) {
                    let distance = Double(abs(iPos - jPos))
                    minDistance = min(minDistance, distance)
                }
            }
        }

        // Convert distance to score (closer is better)
        return minDistance < Double.infinity ? 1.0 / (1.0 + minDistance) : 0.0
    }

    // MARK: - Real LLM Generation with Gemini

    func generate(context: [String], query: String) async throws -> (String, String) {
        let contextText = context.joined(separator: "\n\n")
        
        let prompt = """
        You are an expert assistant that provides concise and accurate answers based on the given context.

        Context:
        \(contextText)

        Question: \(query)

        Please provide a clear, concise answer based only on the information provided in the context. If the context doesn't contain enough information to answer the question, say so. Keep your response focused and relevant.
        """
        
        let response = try await geminiService.generate(prompt: prompt, query: query)
        return (prompt, response)
    }

    @MainActor
    func runRAGPipeline(query: String) -> AsyncStream<RAGDemoResult> {
        return AsyncStream { continuation in
            Task {
                var result = RAGDemoResult()

                do {
                    guard let docURL = Bundle.main.url(forResource: "demo", withExtension: "md") else {
                        throw NSError(domain: "demo.md not found in bundle", code: 1)
                    }
                    let documentText = try loadDocument(docURL)

                    // Slicing
                    let slicesResult = slice(text: documentText)
                    result.slices = slicesResult
                    continuation.yield(result)

                    // Indexing (now synchronous)
                    let indexedResult = indexWithEmbeddings(slices: slicesResult)
                    result.indexed = indexedResult
                    continuation.yield(result)

                    // Retrieval (now synchronous)
                    let retrievedResult = retrieve(index: indexedResult, query: query)
                    result.retrieved = retrievedResult
                    continuation.yield(result)

                    // Reranking
                    let rerankedResult = rerank(chunks: retrievedResult, query: query)
                    result.reranked = rerankedResult
                    continuation.yield(result)

                    // Generation (now async with real LLM)
                    let generatedResult = try await generate(context: rerankedResult, query: query)
                    result.generated = generatedResult
                    continuation.yield(result)

                    continuation.finish()
                } catch {
                    result.error = error
                    continuation.yield(result)
                    continuation.finish()
                }
            }
        }
    }
}
