import torch
from dataclasses import dataclass
from typing import Callable, Optional
from src.config import DeviceConfig
from src.logger import get_logger
from vllm import LLM, SamplingParams
from vllm.transformers_utils.tokenizer import get_tokenizer


logger = get_logger(__name__)


@dataclass
class ModelConfig(DeviceConfig):
    model_id: str = ""
    max_tokens: Optional[int] = None
    max_new_tokens: Optional[int] = None
    temperature: float = 0
    top_p: float = 0.95
    prompt_template: str = "{text}"
    system_prompt: Optional[str] = None

    @classmethod
    def from_yaml(cls, path: str) -> "ModelConfig":
        import yaml

        with open(path, "r") as file:
            config_data = yaml.safe_load(file)
        return cls(**config_data)


def load_model(config: ModelConfig):
    assert config.model_id, "Model ID must be specified in the configuration."

    model = LLM(
        model=config.model_id,
        tensor_parallel_size=config.num_gpus,
        trust_remote_code=True,
        max_model_len=config.max_tokens,
        dtype=torch.bfloat16,
        distributed_executor_backend="mp",
        gpu_memory_utilization=config.gpu_memory_utilization,
        enforce_eager=config.enforce_eager,
    )

    tokenizer = get_tokenizer(config.model_id, tokenizer_mode="auto", trust_remote_code=True)
    terminators = [
        tokenizer.eos_token_id,
    ]

    tokenizer.eot_token_id = terminators

    sampling_params = SamplingParams(
        max_tokens=config.max_new_tokens,
        temperature=config.temperature,
        top_p=config.top_p,
    )

    return (model, tokenizer, sampling_params)


class LLMGenerator:
    def __init__(self, config: ModelConfig) -> None:
        self.model, self.tokenizer, self.sampling_params = load_model(config)
        self.system_prompt = config.system_prompt

    def generate(self, samples: list, prompt: Callable[[str], str]) -> list[str]:
        def build_conversation(sample: str) -> list[dict]:
            conversation = []
            if self.system_prompt:
                conversation.append({"role": "system", "content": self.system_prompt})
            conversation.append({"role": "user", "content": prompt(sample)})
            return conversation

        messages = [build_conversation(sample) for sample in samples]
        prompts = [
            self.tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True)
            for message in messages
        ]

        responses = self.model.generate(prompts, self.sampling_params, use_tqdm=True)

        for i, response in enumerate(responses):
            if response.outputs[0].finish_reason == "length":
                logger.warning(
                    f"Response {i}/{len(samples)} truncated due to max token limit. Response:\n{response.outputs[0].text}\n"
                )

        return [output.outputs[0].text for output in responses]

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        del self.model

        import gc

        torch.cuda.empty_cache()
        gc.collect()
