import os
import faiss
import json
from sentence_transformers import SentenceTransformer

KNOWLEDGE_BASE_DIR = "./knowledge_base"
FAISS_INDEX_PATH = "./faiss_index.bin"
METADATA_PATH = "./faiss_metadata.json"
EMBEDDING_MODEL = "BAAI/bge-small-en-v1.5"

print(f"--- Loading Embedding Model: {EMBEDDING_MODEL} ---")
embedder = SentenceTransformer(EMBEDDING_MODEL)
text_chunks = []

def chunk_document(text):
    paragraphs = text.split("\n\n")
    return [p.strip() for p in paragraphs if len(p.strip()) > 20]

def initialize_vector_store():
    global text_chunks
    embedding_dim = embedder.get_sentence_embedding_dimension()
    
    if os.path.exists(FAISS_INDEX_PATH) and os.path.exists(METADATA_PATH):
        print("--- Loading existing FAISS index ---")
        index = faiss.read_index(FAISS_INDEX_PATH)
        with open(METADATA_PATH, "r") as f:
            text_chunks = json.load(f)
        return index

    print("--- Building FAISS index from scratch ---")
    os.makedirs(KNOWLEDGE_BASE_DIR, exist_ok=True)
    
    all_text = ""
    for filename in os.listdir(KNOWLEDGE_BASE_DIR):
        if filename.endswith(".txt") or filename.endswith(".md"):
            with open(os.path.join(KNOWLEDGE_BASE_DIR, filename), "r", encoding="utf-8") as f:
                all_text += f.read() + "\n\n"
                
    if not all_text.strip():
        all_text = "Flow matching is a generative modeling framework."

    text_chunks = chunk_document(all_text)
    
    print(f"--- Embedding {len(text_chunks)} chunks ---")
    embeddings = embedder.encode(text_chunks, convert_to_numpy=True)
    
    index = faiss.IndexFlatL2(embedding_dim)
    index.add(embeddings)
    
    faiss.write_index(index, FAISS_INDEX_PATH)
    with open(METADATA_PATH, "w") as f:
        json.dump(text_chunks, f)
        
    return index

def retrieve_context(index, query, top_k=3):
    query_vector = embedder.encode([query], convert_to_numpy=True)
    distances, indices = index.search(query_vector, k=top_k)
    retrieved_chunks = [text_chunks[i] for i in indices[0] if i != -1]
    return "\n\n---\n\n".join(retrieved_chunks)