from typing import Dict, Any
import logging

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftConfig, PeftModel
import torch.cuda


LOGGER = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
device = "cuda" if torch.cuda.is_available() else "cpu"


class EndpointHandler():
    def __init__(self, path=""):
        config = PeftConfig.from_pretrained(path)
        model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, load_in_8bit=True, trust_remote_code=True, device_map='auto')
        self.tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
        self.tokenizer.chat_template = "{% for message in messages %}{% if message['role'] == 'user' %}{{ bos_token + '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'system' %}{{ bos_token + '<<SYS>>\\n' + message['content'] + '\\n<</SYS>>\\n\\n' }}{% elif message['role'] == 'assistant' %}{{ ' '  + message['content'] + ' ' + eos_token }}{% endif %}{% endfor %}"
        # Load the Lora model
        self.model = PeftModel.from_pretrained(model, path)

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        LOGGER.info(f"Received data: {data}")
        # Get inputs
        # Extract required parameters from data
        message = data.get("message")
        chat_history = data.get("chat_history", [])
        system_prompt = data.get("system_prompt", "")

        # Extract optional parameters for the generate function logic
        instruction = data.get("instruction")
        conclusions = data.get("conclusions")
        context = data.get("context")

        # Optional parameters with default values
        max_new_tokens = data.get("max_new_tokens", 1024)
        temperature = data.get("temperature", 0.6)
        top_p = data.get("top_p", 0.9)
        top_k = data.get("top_k", 50)
        repetition_penalty = data.get("repetition_penalty", 1.2)

        if message is None or system_prompt is None:
            raise ValueError("Missing required parameters.")

        # Call the generate function
        output = generate(
            tokenizer=self.tokenizer,
            model=self.model,
            message=message,
            chat_history=chat_history,
            system_prompt=system_prompt,
            instruction=instruction,
            conclusions=conclusions,
            context=context,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            repetition_penalty=repetition_penalty
        )

        # Postprocess
        prediction = output
        LOGGER.info(f"Generated text: {prediction}")
        return {"generated_text": prediction}

def generate(
    tokenizer,
    model,
    message: str,
    chat_history: list[tuple[str, str]],
    system_prompt: str = None,
    instruction: str = None,
    conclusions: list[tuple[str, str]] = None,
    context: list[str] = None,
    max_new_tokens: int = 1024,
    temperature: float = 0.6,
    top_p: float = 0.9,
    top_k: int = 50,
    repetition_penalty: float = 1.2,
) -> str:
    
    # Check if the system_prompt is provided, else construct it from instruction, conclusions, and context
    if system_prompt is None and instruction is not None and conclusions is not None and context is not None:
        system_prompt = "Instruction: {}\nConclusions:\n".format(instruction)
        for idx, (conclusion, conclusion_key) in enumerate(conclusions):
            system_prompt += "{}: {}\n".format(conclusion_key, conclusion)
        system_prompt += "\nContext:\n"
        for idx, ctx in enumerate(context):
            system_prompt += "{}: [{}]\n".format(ctx, idx + 1)

    # Construct conversation history
    conversation = [{"role": "system", "content": system_prompt}]
    for user, assistant in chat_history:
        if user is not None:
            conversation.extend([{"role": "user", "content": user}])
        conversation.extend([{"role": "assistant", "content": assistant}])
    conversation.append({"role": "user", "content": message})

    # Tokenize and process the conversation
    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
    input_ids = input_ids.to(model.device)

    # Generate the response
    return tokenizer.decode(model.generate(input_ids.input_ids, max_new_tokens=max_new_tokens))