# -----------------------------------------------
# build_faiss.py
# Reads medical_knowledge.txt
# Splits into chunks
# Embeds with sentence-transformers
# Saves FAISS index to faiss_index/
# Run once after prepare_data.py
# -----------------------------------------------

from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
from config import (
    MEDICAL_KNOWLEDGE_FILE,
    FAISS_INDEX_PATH,
    CHUNK_SIZE,
    CHUNK_OVERLAP,
    EMBEDDING_MODEL
)
import os

# -----------------------------------------------
# Load knowledge base from txt file
# -----------------------------------------------
def load_knowledge_base():
    print(f"Loading {MEDICAL_KNOWLEDGE_FILE}...")

    with open(MEDICAL_KNOWLEDGE_FILE, "r", encoding="utf-8") as f:
        content = f.read()

    # Split on separator
    raw_entries = content.split("\n---\n")

    # Convert to LangChain Documents
    docs = []
    for entry in raw_entries:
        entry = entry.strip()
        if not entry:
            continue

        # Extract disease name for metadata
        disease = "Unknown"
        for line in entry.split("\n"):
            if line.startswith("DISEASE:"):
                disease = line.replace("DISEASE:", "").strip()
                break

        docs.append(Document(
            page_content=entry,
            metadata={"disease": disease}
        ))

    print(f"✅ Loaded {len(docs)} documents")
    return docs

# -----------------------------------------------
# Split documents into chunks
# -----------------------------------------------
def split_documents(docs):
    print(f"\nSplitting into chunks...")
    print(f"  Chunk size: {CHUNK_SIZE} chars")
    print(f"  Overlap: {CHUNK_OVERLAP} chars")

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
        separators=["\n---\n", "\n\n", "\n", " "]
    )

    chunks = splitter.split_documents(docs)
    print(f"✅ Created {len(chunks)} chunks")
    return chunks

# -----------------------------------------------
# Build and save FAISS index
# -----------------------------------------------
def build_faiss(chunks):
    print(f"\nLoading embedding model: {EMBEDDING_MODEL}")
    print("(This downloads ~90MB on first run, then cached forever)")

    embeddings = HuggingFaceEmbeddings(
        model_name=EMBEDDING_MODEL,
        model_kwargs={"device": "cpu"},
        encode_kwargs={"normalize_embeddings": True}
    )
    print("✅ Embedding model loaded")

    print("\nBuilding FAISS index...")
    print("(Embedding all entries — takes 2-5 minutes)")

    vectorstore = FAISS.from_documents(chunks, embeddings)

    os.makedirs(FAISS_INDEX_PATH, exist_ok=True)
    vectorstore.save_local(FAISS_INDEX_PATH)

    print(f"✅ FAISS index saved to {FAISS_INDEX_PATH}/")
    return vectorstore

# -----------------------------------------------
# Verify FAISS works with test query
# -----------------------------------------------
def verify_faiss(vectorstore):
    print("\n=== VERIFICATION — Test Queries ===")

    test_queries = [
        "high fever headache joint pain",
        "chest pain shortness of breath",
        "low blood count fatigue pale skin",
    ]

    for query in test_queries:
        results = vectorstore.similarity_search(query, k=2)
        print(f"\nQuery: '{query}'")
        for i, doc in enumerate(results):
            disease = doc.metadata.get("disease", "Unknown")
            print(f"  Result {i+1}: {disease}")

# -----------------------------------------------
# Main
# -----------------------------------------------
if __name__ == "__main__":
    print("=== Medical Report Analyzer — FAISS Builder ===\n")

    # Step 1 — Load
    docs = load_knowledge_base()

    # Step 2 — Split
    chunks = split_documents(docs)

    # Step 3 — Build + Save
    vectorstore = build_faiss(chunks)

    # Step 4 — Verify
    verify_faiss(vectorstore)

    # Step 5 — Final check
    print("\n=== FINAL CHECK ===")
    index_files = os.listdir(FAISS_INDEX_PATH)
    print(f"✅ Files in faiss_index/: {index_files}")
    print("\n✅ FAISS build complete!")
    print("Next step: Build agents!")