import sys
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."""
    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)
    


def main():
    # 1. Setup the LLM
    print("🚀 Initializing Qwen-8B via LM Studio...")
    lm = dspy.LM("lm_studio/qwen/qwen3-8b", api_base="http://192.168.0.49:1234/v1/")
    dspy.configure(lm=lm)

    # 2. Load the RAG System (only happens once!)
    print("📚 Loading FAISS index and campaign notes...")
    try:
        rag_system = DnDRAG()
        print("✅ Ready! Ask me anything about the campaign. (Type 'exit' or 'q' to quit)")
    except Exception as e:
        print(f"❌ Failed to initialize: {e}")
        return

    # 3. Interactive Loop
    while True:
        try:
            print("\n" + "─" * 30)
            query = input("📝 Query: ").strip()

            # Exit conditions
            if query.lower() in ["exit", "quit", "q"]:
                print("Farewell, traveler. Good luck on your quest!")
                break
            
            if not query:
                continue

            print("🔍 Searching and thinking...")
            response = rag_system(question=query)

            # Print Response
            print("\n📜 AI RESPONSE:")
            print(response.answer)

        except KeyboardInterrupt:
            print("\n\nExiting... See you next session!")
            sys.exit(0)
        except Exception as e:
            print(f"\n⚠️ An error occurred: {e}")

if __name__ == "__main__":
    main()