import faiss
from sentence_transformers import SentenceTransformer
import numpy as np
import os
import pickle
from typing import List, Tuple
from .config import config

class VectorStore:
    def __init__(self):
        print(f"Loading embedding model: {config.EMBEDDING_MODEL}...")
        self.model = SentenceTransformer(config.EMBEDDING_MODEL)
        self.dimension = 384 # all-MiniLM-L6-v2 dimension
        self.index = faiss.IndexFlatL2(self.dimension)
        self.chunks = [] # Keep metadata in memory for this simple version
        self.load_index()

    def add_chunks(self, chunks: List[str], metadatas: List[dict]):
        embeddings = self.model.encode(chunks)
        self.index.add(np.array(embeddings).astype('float32'))
        self.chunks.extend(zip(chunks, metadatas))
        self.save_index()

    def search(self, query: str, k: int = 5, notebook_id: str = None) -> List[Tuple[str, dict, float]]:
        query_vector = self.model.encode([query])
        
        # Over-fetch to allow for filtering
        search_k = k * 10 
        distances, indices = self.index.search(np.array(query_vector).astype('float32'), search_k)
        
        results = []
        count = 0
        for i, idx in enumerate(indices[0]):
            if idx != -1 and idx < len(self.chunks):
                text, meta = self.chunks[idx]
                
                # Filter by notebook_id if provided
                if notebook_id:
                    if meta.get("notebook_id") != notebook_id:
                        continue
                
                results.append((text, meta, float(distances[0][i])))
                count += 1
                if count >= k:
                    break
                    
        return results

    def save_index(self):
        faiss.write_index(self.index, "vector_store.index")
        with open("chunks_meta.pkl", "wb") as f:
            pickle.dump(self.chunks, f)

    def load_index(self):
        if os.path.exists("vector_store.index") and os.path.exists("chunks_meta.pkl"):
            self.index = faiss.read_index("vector_store.index")
            with open("chunks_meta.pkl", "rb") as f:
                self.chunks = pickle.load(f)

vector_store = VectorStore()