"""
model.py - Inference & Model Orchestration Layer
This module encapsulates the logic for loading and executing high-parameter LLMs. 
It manages the transition from raw user input to structured data by coordinating 
the chat model, embedding model, and Pydantic-based output validation.
"""

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.output_parsers import PydanticOutputParser
from schema import PokemonEmeraldResponse
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

CHAT_MODEL_ID = "mistralai/Mistral-Small-Instruct-2409"
EMBEDDING_MODEL_ID = "nomic-ai/nomic-embed-text-v1.5"

def load_chat_model():
    """
    Bootstrap the tokenizer and chat model for the RAG pipeline.
    Returns: (model, tokenizer) tuple.
    """
    print(f"Loading {CHAT_MODEL_ID} model...")
    tokenizer = AutoTokenizer.from_pretrained(CHAT_MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(
        CHAT_MODEL_ID,
        torch_dtype="auto",
        device_map="auto"
    )
    print("Model loaded!")
    return model, tokenizer

def load_embedding_model():
    """
    Bootstrap the HuggingFaceEmbeddings model for the RAG pipeline.
    Returns: An instance of HuggingFaceEmbeddings configured with EMBEDDING_MODEL_ID.
    """
    print(f"Loading {EMBEDDING_MODEL_ID} model...")
    embedding_model = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_ID)
    print("Model loaded!")
    return embedding_model
    
def generate_answer(chat_model, tokenizer, prompt, max_retries=3):
    """
    Executes the core RAG generation loop with automated schema validation.
    
    This function manages the transition from raw model output to structured JSON. 
    It includes a retry mechanism to handle edge-case parsing failures, ensuring 
    the final response adheres to the PokemonEmeraldResponse Pydantic schema.
    """
    print("Generating chat response...")
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    inputs_size = len(inputs.input_ids[0])
    parser = PydanticOutputParser(pydantic_object=PokemonEmeraldResponse)
    attempt = 0
    while attempt < max_retries:
        try:
            response = chat_model.generate(
                **inputs, 
                max_new_tokens=256, 
                num_beams=3,
                early_stopping=True,
                repetition_penalty=1.2,
                length_penalty=1,
            )
            output = tokenizer.decode(response[0][inputs_size:], skip_special_tokens=True)
            return parser.parse(output)
        except Exception as e:
            print(f"Chat model response generation failed, attempt {attempt}.")
            attempt += 1
    print(f"Parsing failed. Output: {output}")