import torch

from typing import Dict, List, Any
from transformers import AutoTokenizer, BitsAndBytesConfig
from peft import AutoPeftModelForCausalLM


def parse_output(text):
    marker = "### Response:"
    if marker in text:
        pos = text.find(marker) + len(marker)
    else:
        pos = 0
    return text[pos:].replace("<pad>", "").replace("</s>", "").strip()


class EndpointHandler:
    def __init__(self, path="./", use_bnb=True):

        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
        )
        self.model = AutoPeftModelForCausalLM.from_pretrained(
            path, load_in_8bit=False, quantization_config=bnb_config, device_map="auto"
        )

        self.tokenizer = AutoTokenizer.from_pretrained(path)
        print("Memory footprint: ", self.model.get_memory_footprint())
        print("Device map: ", self.model.hf_device_map)

    def __call__(self, data: Any) -> List[List[Dict[str, str]]]:


        inputs = data.get("inputs", data)
        prompt = f"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction: \n{inputs}\n\n### Response: \n"
        parameters = data.get("parameters", {})
        with torch.no_grad():
            inputs = self.tokenizer(
                prompt, return_tensors="pt", return_token_type_ids=False
            ).to(self.model.device)
            outputs = self.model.generate(**inputs, **parameters)

        return {
            "generated_text": parse_output(
                self.tokenizer.decode(outputs[0].tolist(), skip_special_tokens=True)
            )
        }