from typing import List, Dict
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

class EndpointHandler:
    def __init__(self, path: str):
        # Load model and tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.model = AutoModelForCausalLM.from_pretrained(
            path,
            torch_dtype=torch.float32,  # Use float32 for CPU
            device_map="auto"
        )
        
        # Set up generation parameters
        self.default_params = {
            "max_length": 1000,
            "temperature": 0.7,
            "top_p": 0.7,
            "top_k": 50,
            "repetition_penalty": 1.0,
            "do_sample": True,
            "pad_token_id": self.tokenizer.pad_token_id,
            "eos_token_id": self.tokenizer.eos_token_id
        }

    def __call__(self, data: Dict):
        try:
            # Handle input
            if isinstance(data.get("inputs"), str):
                input_text = data["inputs"]
            else:
                input_text = data.get("inputs")[0] if isinstance(data.get("inputs"), list) else str(data.get("inputs"))

            # Print debug information
            print(f"Input text: {input_text}")

            # Tokenize with fixed dimensions
            tokenizer_output = self.tokenizer(
                input_text,
                padding='max_length',  # Changed to max_length
                truncation=True,
                max_length=512,  # Fixed length
                return_tensors="pt",
                return_attention_mask=True
            )

            # Print tensor shapes for debugging
            print(f"Input ids shape: {tokenizer_output['input_ids'].shape}")
            print(f"Attention mask shape: {tokenizer_output['attention_mask'].shape}")

            # Generate response
            with torch.no_grad():
                outputs = self.model.generate(
                    tokenizer_output["input_ids"],
                    attention_mask=tokenizer_output["attention_mask"],
                    max_new_tokens=1024,
                    pad_token_id=self.tokenizer.pad_token_id,
                    do_sample=True,
                    temperature=0.7,
                    top_p=0.7,
                    top_k=50
                )

            # Decode response
            generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            
            return [{"generated_text": generated_text}]
            
        except Exception as e:
            print(f"Error in generation: {str(e)}")
            print(f"Model config: {self.model.config}")
            return {"error": str(e)}

    def preprocess(self, request):
        """
        Prepare request for inference
        """
        if request.content_type != "application/json":
            raise ValueError("Content type must be application/json")
        
        data = request.json
        return data

    def postprocess(self, data):
        """
        Post-process model output
        """
        return data