10. Sentence Embeddings & Semantic Search

// Implementasi semantic search engine dengan NLEmbedding

import NaturalLanguage
import Foundation

// Vektor yang disimpan di-disk (simple implementation)
struct EmbeddedDocument: Codable {
    let id: String
    let text: String
    let vector: [Double]
    let metadata: [String: String]
}

final class SemanticDocumentStore {
    private var documents: [EmbeddedDocument] = []
    private let embedding: NLEmbedding
    private let storeURL: URL

    init?(storeURL: URL) {
        guard let embedding = NLEmbedding.sentenceEmbedding(for: .english) else { return nil }
        self.embedding = embedding
        self.storeURL = storeURL
        loadFromDisk()
    }

    // Tambah dokumen — hitung embedding saat insert
    func addDocument(id: String, text: String, metadata: [String: String] = [:]) {
        guard let vector = embedding.vector(for: text) else { return }

        let doc = EmbeddedDocument(id: id, text: text, vector: vector, metadata: metadata)
        documents.removeAll { $0.id == id }  // Upsert
        documents.append(doc)
        saveToDisk()
    }

    // Semantic search
    func search(query: String, topK: Int = 10, threshold: Double = 0.5) -> [SearchResult] {
        guard let queryVector = embedding.vector(for: query) else { return [] }

        return documents
            .compactMap { doc -> SearchResult? in
                let score = cosineSimilarity(queryVector, doc.vector)
                guard score >= threshold else { return nil }
                return SearchResult(document: doc, score: score)
            }
            .sorted { $0.score > $1.score }
            .prefix(topK)
            .map { $0 }
    }

    private func cosineSimilarity(_ v1: [Double], _ v2: [Double]) -> Double {
        guard v1.count == v2.count else { return 0 }
        let dot = zip(v1, v2).reduce(0.0) { $0 + $1.0 * $1.1 }
        let mag1 = sqrt(v1.reduce(0.0) { $0 + $1 * $1 })
        let mag2 = sqrt(v2.reduce(0.0) { $0 + $1 * $1 })
        guard mag1 > 0 && mag2 > 0 else { return 0 }
        return dot / (mag1 * mag2)
    }

    private func saveToDisk() {
        guard let data = try? JSONEncoder().encode(documents) else { return }
        try? data.write(to: storeURL)
    }

    private func loadFromDisk() {
        guard let data = try? Data(contentsOf: storeURL),
              let loaded = try? JSONDecoder().decode([EmbeddedDocument].self, from: data) else { return }
        documents = loaded
    }
}

struct SearchResult {
    let document: EmbeddedDocument
    let score: Double
}