from typing import Dict, List, Any
import torch
from transformers import LlavaNextVideoForConditionalGeneration, LlavaNextVideoProcessor
from peft import PeftModel

class EndpointHandler:
    def __init__(self):
        self.base_model_name = "llava-hf/LLaVA-NeXT-Video-7B-hf"
        self.adapter_model_name = "EnariGmbH/surftown-1.0"

        # Load the base model
        self.model = LlavaNextVideoForConditionalGeneration.from_pretrained(
            self.base_model_name,
            torch_dtype=torch.float16,
            device_map="auto"
        )

        # Load the fine-tuned adapter model into the base model
        self.model = PeftModel.from_pretrained(self.model, self.adapter_model_name)

        # Merge the adapter weights into the base model and unload the adapter
        self.model = self.model.merge_and_unload()

        # # Optionally, load and save the processor (if needed)
        self.processor = LlavaNextVideoProcessor.from_pretrained(self.adapter_model_name)

        # Ensure the model is in evaluation mode
        self.model.eval()

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Args:
            data (Dict): Contains the input data including "clip" and "prompt".
        
        Returns:
            List[Dict[str, Any]]: The generated text from the model.
        """
        # Extract inputs from the data dictionary
        clip = data.get("clip")
        prompt = data.get("prompt")

        if clip is None or prompt is None:
            return [{"error": "Missing 'clip' or 'prompt' in input data"}]

        # Prepare the inputs for the model
        inputs_video = self.processor(text=prompt, videos=clip, padding=True, return_tensors="pt").to(self.model.device)

        # Generate output from the model
        generate_kwargs = {"max_new_tokens": 512, "do_sample": True, "top_p": 0.9}
        output = self.model.generate(**inputs_video, **generate_kwargs)
        generated_text = self.processor.batch_decode(output, skip_special_tokens=True)

        # Extract the relevant part of the assistant's answer
        assistant_answer_start = generated_text[0].find("ASSISTANT:") + len("ASSISTANT:")
        assistant_answer = generated_text[0][assistant_answer_start:].strip()

        return [{"generated_text": assistant_answer}]