//
//  VectorManager.swift
//  vectorsearch
//
//  Created by Du Qingjie on 2024/09/06.
//

import SimilaritySearchKit
import Foundation

class VectorManager: ObservableObject {
    static let shared = VectorManager()

    private var similarityIndex: SimilarityIndex?
    
    init() {
        Task {
            await load()
        }
    }

    func load() async {
        similarityIndex = await SimilarityIndex(
            model: NativeEmbeddings(),
            metric: CosineSimilarity()
        )
        if let data = UserDefaults.standard.data(forKey: "vector_data") {
            let items = try? JSONDecoder().decode([SimilarityIndex.IndexItem].self, from: data)
            similarityIndex?.indexItems = items ?? []
        }else{
            similarityIndex?.indexItems = []
        }
    }
    
    func setFileItem(path: String, text: String, modify_time: Int) async{
        guard let index = similarityIndex else { return }
        
        let curItem = index.getItem(id: path)
        if(curItem != nil && Int(curItem?.metadata["mt"] ?? String(Int.max)) ?? Int.max > Int(modify_time)){
            index.updateItem(id: path, text: text, embedding: nil, metadata: ["mt": String(modify_time)] )
            if let encoded = try? JSONEncoder().encode(index.indexItems) {
                UserDefaults.standard.set(encoded, forKey: "vector_data")
            }
        }else if(curItem == nil){
            await index.addItem(id: path, text: text, metadata: ["mt": String(modify_time)])
            if let encoded = try? JSONEncoder().encode(index.indexItems) {
                UserDefaults.standard.set(encoded, forKey: "vector_data")
            }
        }
    }
    
    func search(query: String) async -> [SimilarityIndex.SearchResult]{
        guard let index = similarityIndex else { return []}
        
        let results = await index.search(query, top: 100)
        return results
    }

    func deleteFileItem(path: String){
        guard let index = similarityIndex else { return }
        
        index.removeItem(id: path)
        UserDefaults.standard.setValue(index.indexItems, forKey: "vector_data")
    }

    func clear() {
        guard let index = similarityIndex else { return }
        
        index.removeAll()
        UserDefaults.standard.setValue(index.indexItems, forKey: "vector_data")
    }

    func getFileItem(id: String) -> SimilarityIndex.IndexItem?{
        guard let index = similarityIndex else { return nil }
        
        return index.getItem(id: id)
    }
}




