import dspy
from langchain_community.vectorstores import FAISS
from embedding import LocalLMEmbeddings
from pathlib import Path

# --- DSPy Signature ---
class DnDContextQA(dspy.Signature):
    """Answer DnD campaign questions using provided snippets and full file context.
    /no_think"""
    context = dspy.InputField(desc="Relevant chunks and full file contents from the campaign notes.")
    question = dspy.InputField()
    answer = dspy.OutputField(desc="A detailed answer based on the notes, citing the source file.")

# --- DSPy Module ---
class DnDRAG(dspy.Module):
    def __init__(self, db_path="./local_faiss_db", k=3):
        super().__init__()
        # 1. Setup Embeddings & Load FAISS
        self.embeddings = LocalLMEmbeddings(
            model="text-embedding-qwen3-embedding-8b", 
            base_url="http://192.168.0.49:1234"
        )
        self.vectorstore = FAISS.load_local(
            db_path, self.embeddings, allow_dangerous_deserialization=True
        )
        self.k = k
        
        # 2. Setup the Predictor (Chain of Thought for better reasoning)
        self.generate_answer = dspy.ChainOfThought(DnDContextQA)

    def get_full_file_content(self, file_path):
        """Helper to read the full source file if it exists."""
        try:
            return Path(file_path).read_text(encoding='utf-8')
        except Exception:
            return ""

    def forward(self, question):
        # 1. Search for top-k chunks
        results = self.vectorstore.similarity_search(question, k=self.k)
        
        # 2. Extract unique file paths to load "Full Context"
        # This prevents the LLM from being 'blind' to the rest of a relevant session note
        unique_paths = list(set([doc.metadata.get("full_path") for doc in results]))
        
        context_parts = []
        for i, doc in enumerate(results):
            source = doc.metadata.get("source", "Unknown")
            context_parts.append(f"--- Chunk {i+1} from {source} ---\n{doc.page_content}")

        # 3. Add the Full Content of the top match (optional, but requested!)
        # We'll just take the top 1 file to avoid context window explosion
        if unique_paths:
            top_file_content = self.get_full_file_content(unique_paths[0])
            context_parts.append(f"\n=== FULL SOURCE FILE: {Path(unique_paths[0]).name} ===\n{top_file_content[:10000]}")

        # 4. Join everything into one context string
        context_str = "\n\n".join(context_parts)
        
        # 5. Generate Response
        prediction = self.generate_answer(context=context_str, question=question)
        return dspy.Prediction(answer=prediction.answer, context=context_str)