# handler.py
"""
Hugging Face Inference Endpoint Handler for SciGuru PPO+LoRA Model
Handles inference requests for the fine-tuned Llama model with LoRA adapters
"""

import torch
from typing import Dict, List, Any
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig
)
from peft import PeftModel
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class EndpointHandler:
    def __init__(self, path=""):
        """
        Initialize the handler by loading the model and tokenizer.
        
        Args:
            path: Path to the model directory containing the LoRA adapters
        """
        logger.info(f"Initializing EndpointHandler with path: {path}")
        
        # Model configuration - must match training setup
        self.base_model_name = "meta-llama/Llama-3.2-3B-Instruct"
        self.max_seq_length = 512
        
        # Load tokenizer
        logger.info("Loading tokenizer...")
        self.tokenizer = AutoTokenizer.from_pretrained(self.base_model_name)
        
        # Set up tokenizer properly
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        self.tokenizer.padding_side = "right"
        
        # Quantization config for memory efficiency (same as training)
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4"
        )
        
        # Load base model with quantization
        logger.info("Loading base model with 4-bit quantization...")
        self.model = AutoModelForCausalLM.from_pretrained(
            self.base_model_name,
            quantization_config=bnb_config,
            device_map="auto",
            torch_dtype=torch.float16,
            trust_remote_code=True
        )
        
        # Load LoRA adapters
        logger.info(f"Loading LoRA adapters from {path}...")
        self.model = PeftModel.from_pretrained(
            self.model,
            path,
            torch_dtype=torch.float16
        )
        
        # Set model to evaluation mode
        self.model.eval()
        
        # Enable memory efficient operations
        if torch.cuda.is_available():
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True
        
        logger.info("Model initialization complete!")
        
    def format_prompt(self, question: str) -> str:
        """
        Format the question with the same prompt template used during training.
        
        Args:
            question: The scientific question to answer
            
        Returns:
            Formatted prompt string
        """
        prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
You are an expert science educator. Explain scientific concepts clearly and simply, using analogies and everyday examples when possible. Make your explanations accessible to someone with basic high school science knowledge.<|eot_id|>
<|start_header_id|>user<|end_header_id|>
{question}<|eot_id|>
<|start_header_id|>assistant<|end_header_id|>
"""
        return prompt
    
    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Handle inference requests.
        
        Args:
            data: Input data containing 'inputs' and optional 'parameters'
                - inputs: Can be a string or list of strings (questions)
                - parameters: Optional generation parameters
                
        Returns:
            List of dictionaries containing generated explanations
        """
        # Extract inputs
        inputs = data.get("inputs", "")
        parameters = data.get("parameters", {})
        
        # Handle single string or list of strings
        if isinstance(inputs, str):
            questions = [inputs]
        elif isinstance(inputs, list):
            questions = inputs
        else:
            raise ValueError("Inputs must be a string or list of strings")
        
        # Extract generation parameters with defaults matching training
        max_new_tokens = parameters.get("max_new_tokens", 256)
        temperature = parameters.get("temperature", 0.7)
        top_p = parameters.get("top_p", 0.9)
        do_sample = parameters.get("do_sample", True)
        num_return_sequences = parameters.get("num_return_sequences", 1)
        
        # Process each question
        results = []
        
        for question in questions:
            try:
                # Format the prompt
                formatted_prompt = self.format_prompt(question)
                
                # Tokenize
                inputs_encoded = self.tokenizer(
                    formatted_prompt,
                    return_tensors="pt",
                    truncation=True,
                    max_length=self.max_seq_length // 2,  # Leave room for response
                    padding=False
                ).to(self.model.device)
                
                prompt_length = inputs_encoded.input_ids.shape[1]
                
                # Generate response
                with torch.no_grad():
                    outputs = self.model.generate(
                        inputs_encoded.input_ids,
                        attention_mask=inputs_encoded.attention_mask,
                        max_new_tokens=max_new_tokens,
                        do_sample=do_sample,
                        temperature=temperature,
                        top_p=top_p,
                        pad_token_id=self.tokenizer.pad_token_id,
                        eos_token_id=self.tokenizer.eos_token_id,
                        num_return_sequences=num_return_sequences
                    )
                
                # Decode responses
                generated_texts = []
                for output in outputs:
                    # Extract only the generated part (after the prompt)
                    generated_ids = output[prompt_length:]
                    generated_text = self.tokenizer.decode(
                        generated_ids,
                        skip_special_tokens=True,
                        clean_up_tokenization_spaces=True
                    )
                    generated_texts.append(generated_text.strip())
                
                # Format result
                if num_return_sequences == 1:
                    results.append({
                        "question": question,
                        "explanation": generated_texts[0]
                    })
                else:
                    results.append({
                        "question": question,
                        "explanations": generated_texts
                    })
                    
            except Exception as e:
                logger.error(f"Error processing question '{question}': {str(e)}")
                results.append({
                    "question": question,
                    "error": str(e)
                })
        
        # Clear cache to free memory
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        return results