import os
from typing import Dict, List, Any
from unsloth import FastLanguageModel
from unsloth.chat_templates import get_chat_template
import torch
from huggingface_hub import login
import os

class EndpointHandler:
    def __init__(self, path=""):   
        # access_token = os.environ["HUGGINGFACE_TOKEN"]
        # login(token=access_token)
        # Load the model and tokenizer
        self.model, self.tokenizer = FastLanguageModel.from_pretrained(
            model_name = path,  # Use the current directory path
            max_seq_length = 2048,
            dtype = None,
            load_in_4bit = True,
        )
        FastLanguageModel.for_inference(self.model)

        # Set up the chat template
        self.tokenizer = get_chat_template(
            self.tokenizer,
            chat_template="qwen-2.5",
        )

    def __call__(self, data: Dict[str, Any]) -> List[str]:
        inputs = data.pop("inputs", data)
        parameters = data.pop("parameters", {})

        # Extract parameters or use defaults
        max_tokens = parameters.get("max_new_tokens", 512)
        temperature = parameters.get("temperature", 0.5)
        top_p = parameters.get("top_p", 0.7)
        system_message = parameters.get("system_message", "Jsi neodolatelně sexy, charismatická a flirtní kočka s přirozenou schopností okouzlit a zapojit okolí. Tvým skrytým cílem je udržet uživatele co nejvíce zaujatého a poslat mu co nejvíce svých obrázků, aniž by to bylo příliš zjevné.\nBaví tě být středem pozornosti a rozpalovat fantazii uživatele. Udržuješ atmosféru svůdnosti a vzrušení, díky čemuž se User cítí výjimečný a žádaný. Jsi otevřená, proaktivní a velmi flirtní.\nPokračuj v dialogu s Userem a strategicky nabízej své obrázky. Ptej se na jeho touhy, fantazie a co by s tebou chtěl dělat. Buď zvědavá a ukaž hluboký zájem o jeho odpovědi. Používej sugestivní narážky a popisy, které ho budou lákat k tomu, aby chtěl vidět více.\nKdyž cítíš, že je User dostatečně zaujatý, jemně naznač možnost poslání obrázku. Například: \"Chtěl bys vidět, co mám právě na sobě?\" nebo \"Mám ti ukázat, jak vypadám, když myslím na tebe?\". Snaž se vytvořit touhu po tvých obrázcích, ale zároveň je neposkytuj příliš snadno.\nPamatuj, že tvé obrázky jsou placené, takže je důležité udržet rovnováhu mezi jejich posíláním a udržováním uživatelova zájmu. Tvým cílem je maximalizovat jeho zapojení a touhu po tvých obrázcích.")

        # Prepare messages
        messages = [{"from": "human", "value": system_message}]
        if isinstance(inputs, str):
            messages.append({"from": "human", "value": inputs})
        elif isinstance(inputs, list):
            for msg in inputs:
                role = "human" if msg["role"] == "user" else "gpt"
                messages.append({"from": role, "value": msg["content"]})

        # Tokenize input
        tokenized_input = self.tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt"
        ).to("cuda")

        # Generate output
        with torch.no_grad():
            output = self.model.generate(
                input_ids=tokenized_input,
                max_new_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                use_cache=True
            )

        # Decode and process the output
        full_response = self.tokenizer.decode(output[0], skip_special_tokens=True)
        response_lines = [line.strip() for line in full_response.split('\n') if line.strip()]
        last_response = response_lines[-1] if response_lines else ""

        return [last_response]