"""
helpers.py - RAG Orchestration & Document Processing
This module defines the utility functions for the Pokémon Emerald knowledge base.
It manages the end-to-end retrieval pipeline, including:
- Document Ingestion: Recursive character splitting optimized for technical manuals.
- Vector Management: Integration with ChromaDB for semantic search.
- Prompt Engineering: Context-aware prompt construction with strict Pydantic 
  formatting instructions to ensure structured JSON output.
"""
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import Chroma
from langchain_core.output_parsers import PydanticOutputParser
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from model import load_embedding_model
from schema import PokemonEmeraldResponse

MANUAL_PATH = "emerald_manual.txt"

def load_vector_store(filename):
    """
    Load, split, and index the game manual into a ChromaDB vector store.
    Args:
        filename: Path to the source .txt manual.
    Returns:
        Initialised vector store for similarity search.
    """
    print("Loading vector store...")
    loader = TextLoader(MANUAL_PATH)
    data = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
    chunks = text_splitter.split_documents(data)
    embedding_model = load_embedding_model()
    print(f"Created {len(chunks)} chunks.")
    store = Chroma.from_documents(documents=chunks, embedding=embedding_model)
    return store

def find_similar_documents(vector_store, message, k=5):
    """
    Queries the vector store for manual snippets matching the user's message.
    
    Args:
        vector_store: Initialized Chroma/FAISS instance.
        message: The user's natural language question.
        k: Number of context chunks to retrieve (default=5 for high density).
        
    Returns:
        List of Document objects containing page_content and metadata.
    """
    print("Retrieving relevant documents...")
    context = vector_store.similarity_search(message, k=k) 
    for i, x in enumerate(context):
        print(f"i={i}" + "-"*20)
        print(x)
        print()
    return context

def create_prompt(message, history, context):
    """
    Constructs a grounded, instruction-tuned prompt for the LLM.
    
    Integrates retrieved manual context with a strict system persona to 
    mitigate hallucinations. Injects Pydantic-generated format instructions 
    to ensure the model's response is a valid JSON object matching the 
    PokemonEmeraldResponse schema.
    """
    parser = PydanticOutputParser(pydantic_object=PokemonEmeraldResponse)
    context_text = "\n---\n".join([d.page_content for d in context])
    format_instructions = parser.get_format_instructions()
    prompt = f"""
    System: You are a precise Pokémon Emerald Research Assistant. Your goal is to answer questions using
    ONLY the provided context. If the answer is not contained within the context, state: No relevant
    sources found. Citations should be sufficient to keep context and at least a couple of sentences.
    Include all relevant documents in citations.
    Context: {context_text}
    
    Question: {message}
    {format_instructions}
    Follow the format instructions exactly. The output will be parsed as a single JSON.
    If the answer cannot be found in the provided context, or if you do not know the answer,
    state exactly: "I do not have enough information to answer this."
    Do not repeat previous responses, and do not attempt to guess.
    """
    return prompt