import logging
import re
from typing import List, Optional

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from urartu.common.device import Device
from urartu.common.model import Model
from urartu.utils.dtype import eval_dtype

logger = logging.getLogger(__name__)


class ModelForCausalLM(Model):
    """
    A class for working with a causal language model (CLM) that handles model initialization,
    configuration, and text generation using the model and tokenizer from the Hugging Face
    Transformers library.

    Attributes:
        cfg: An object containing configuration settings such as model name, cache directory,
             API token, and other model-specific settings. Should be a DictConfig for attribute access.

    Methods:
        model: Returns a causal language model with preloaded configurations.
        tokenizer: Returns a tokenizer associated with the causal language model.
        generate: Generates text based on a given prompt and optional generation settings.
    """

    def __init__(self, cfg) -> None:
        """
        Initializes the ModelCausalLanguage class by setting up the base model configuration.

        Args:
            cfg: Configuration object containing necessary parameters like model name,
                 cache directory, device mapping, dtype evaluation, and optional API token.
                 Should be a DictConfig (passed from instantiate).
        """
        super().__init__(cfg)
        self._tokenizer = None

    @property
    def model(self) -> AutoModelForCausalLM:
        """
        Retrieves or instantiates the causal language model specified in the configuration.
        The model is set to evaluation mode with gradients disabled.

        Returns:
            An instance of AutoModelForCausalLM ready for inference.
        """
        if self._model is None:
            # Collect kwargs for from_pretrained
            from_pretrained_kwargs = {
                "cache_dir": self.cfg.get("cache_dir"),
                "device_map": Device.get_device(),
                "token": self.cfg.get("api_token"),
                "trust_remote_code": self.cfg.get("trust_remote_code"),
                "revision": self.cfg.get("revision"),
            }
            # Add dtype if specified
            if self.cfg.get("dtype") is not None:
                from_pretrained_kwargs["torch_dtype"] = eval_dtype(self.cfg.get("dtype"))
            # Add attn_implementation if specified (needed for attention weight extraction)
            if self.cfg.get("attn_implementation") is not None:
                from_pretrained_kwargs["attn_implementation"] = self.cfg.get("attn_implementation")

            self._model = AutoModelForCausalLM.from_pretrained(self.cfg.name, **from_pretrained_kwargs)
            for param in self._model.parameters():
                param.requires_grad = False
            self._model.eval()
        return self._model

    @property
    def tokenizer(self):
        """
        Retrieves or instantiates the tokenizer associated with the causal language model,
        setting the pad token to be the same as the eos token if not already set.

        Returns:
            An instance of AutoTokenizer.
        """
        if self._tokenizer is None:
            self._tokenizer = AutoTokenizer.from_pretrained(self.cfg.name)
            self._tokenizer.pad_token = self._tokenizer.eos_token
        return self._tokenizer

    def generate(self, prompt: str, generate_cfg=None):
        """
        Generates text based on a provided prompt and optional generation settings.

        Args:
            prompt: A string containing the initial text to prompt the model with.
            generate_cfg: An optional dictionary containing generation parameters such as
                          maximum length, number of beams, etc. Defaults to configurations
                          specified in the model's config.

        Returns:
            A string of generated text. If 'output_scores' is included in the generation
            configuration, it returns a tuple of the generated text and the corresponding
            generation scores.
        """
        if not generate_cfg:
            generate_cfg = self.cfg.get("generate")
        self.model.eval()

        prompt_tokenized = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
        prompt_tensor = prompt_tokenized["input_ids"].to(self.model.device)
        attention_mask = prompt_tokenized["attention_mask"].to(self.model.device)
        with torch.no_grad():
            output_tokenized = self.model.generate(input_ids=prompt_tensor, attention_mask=attention_mask, **generate_cfg)
        if "output_scores" in generate_cfg:
            output = self.tokenizer.decode(output_tokenized["sequences"][0], skip_special_tokens=True)
            scores = output_tokenized["scores"]
            if isinstance(scores, (list, tuple)):
                # For single generation, scores are per token, stack them to have shape (num_tokens, vocab_size)
                scores = torch.stack(scores, dim=0)
            return output, scores
        else:
            return self.tokenizer.decode(output_tokenized[0], skip_special_tokens=True)

    def generate_batch(self, prompts: list[str], generate_cfg=None):
        """
        Generates text for a batch of prompts.
        """
        if not generate_cfg:
            generate_cfg = self.cfg.get("generate")
        self.model.eval()

        # Tokenize the batch of prompts
        self.tokenizer.padding_side = "left"  # For decoder-only models
        prompt_tokenized = self.tokenizer(prompts, return_tensors="pt", padding=True, truncation=True)
        prompt_tensor = prompt_tokenized["input_ids"].to(self.model.device)
        attention_mask = prompt_tokenized["attention_mask"].to(self.model.device)

        with torch.no_grad():
            output_tokenized = self.model.generate(input_ids=prompt_tensor, attention_mask=attention_mask, **generate_cfg)

        if "output_scores" in generate_cfg:
            # Decode all sequences in the batch
            outputs = self.tokenizer.batch_decode(output_tokenized.sequences, skip_special_tokens=True)

            # Scores is a tuple of tensors (one for each generated token).
            # Each tensor is of shape (batch_size, vocab_size).
            # We stack them to get a tensor of shape (num_generated_tokens, batch_size, vocab_size).
            scores_stacked = torch.stack(output_tokenized.scores, dim=0)

            # We want to return a list of score tensors, one for each sample in the batch.
            # So we transpose to (batch_size, num_generated_tokens, vocab_size) and split.
            scores_per_sample = list(scores_stacked.transpose(0, 1))

            return outputs, scores_per_sample
        else:
            # Just decode the output sequences
            return self.tokenizer.batch_decode(output_tokenized, skip_special_tokens=True)
