"""ASR model implementation."""

import os

import torch
import torch.nn.functional as F
from transformers import AutoConfig, PreTrainedModel
from transformers.utils import cached_file

from .config import Config


class Model(PreTrainedModel):
    """ASR model that wraps a ``torch.export``-ed model.

    This model wraps a pre-compiled ``torch.export`` artifact (``.pt2``)
    for automatic speech recognition.
    """

    config_class = Config

    def __init__(self, model_name_or_path: str, config: Config):
        super().__init__(config)
        self.model_name_or_path = model_name_or_path
        self.config = config
        self._exported_model = None  # Loaded on demand

    def _load_exported_model(self, device: str):
        """Load the exported model from file."""

        if device not in ["cpu", "cuda"]:
            raise ValueError("Device must be either 'cpu' or 'cuda'.")

        # Determine the path to the exported model file
        filename = f"{self.config.exported_model_file}-{device}.pt2"

        if os.path.isdir(self.model_name_or_path):
            exported_path = os.path.join(self.model_name_or_path, filename)
        else:
            # If it's a model ID from HuggingFace Hub, download the file
            exported_path = cached_file(self.model_name_or_path, filename)

        if not os.path.exists(exported_path):
            raise FileNotFoundError(
                f"Exported model file not found at {exported_path}. "
                f"Make sure the file '{filename}' exists in the model directory."
            )
        self._exported_model = torch.export.load(exported_path).module()

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        """Override from_pretrained to load the exported model without
        requiring standard checkpoint files."""

        # Load config
        config = kwargs.pop("config", None)
        if config is None:
            config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)

        # Create model instance
        model = cls(pretrained_model_name_or_path, config)

        return model

    def compute_mask(self, mask: torch.Tensor) -> torch.Tensor:
        """Compute output mask of model, given input mask.

        Parameters
        ----------
        mask : torch.Tensor
            Input mask of shape `(batch_size, input_steps)` with boolean values.
            True indicates valid positions, False indicates padding.

        Returns
        -------
        torch.Tensor
            Output mask of shape `(batch_size, output_steps)` with boolean values.
            The output will be on the same device as the model.
        """
        if self._exported_model is not None:
            model_device = next(self._exported_model.parameters()).device
        else:
            model_device = torch.device("cpu")

        # Move mask to model device
        mask = mask.to(model_device)

        # Invert mask and convert to float32
        x = (~mask).float()

        # Add channel dimension: (batch_size, input_steps) -> (batch_size, 1, input_steps)
        x = x.unsqueeze(1)

        # Apply max pooling twice with pool_size=4, strides=2, padding=0 (valid)
        for _ in range(2):
            x = F.max_pool1d(x, kernel_size=4, stride=2, padding=0)

        # Invert back and convert to bool
        mask = ~(x.bool())

        # Remove channel dimension: (batch_size, 1, output_steps) -> (batch_size, output_steps)
        mask = mask.squeeze(1)

        return mask

    def forward(self, input_features: torch.Tensor, mask: torch.Tensor = None):
        """Forward pass of the model.

        Parameters
        ----------
        input_features : torch.Tensor
            Input features (e.g., mel-spectrogram features) of shape
            `(batch_size, input_steps, input_features)`).
        mask : torch.Tensor, optional
            Mask of shape `(batch_size, input_steps)` with boolean values.
            True indicates valid positions, False indicates padding.
            If provided, the output will include the computed output mask.

        Returns
        -------
        dict or torch.Tensor
            If mask is provided, returns a dictionary with:
                - 'logits': Model outputs with shape `(batch_size, output_steps, vocab_size)`.
                - 'mask': Output attention mask with shape `(batch_size, output_steps)`.
            Otherwise, returns just the logits tensor.
        """

        if self._exported_model is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
            self._load_exported_model(device)

        # Exported model requires the time dimension to be a multiple of 4.
        batch, input_steps, _ = input_features.shape
        pad = (-input_steps) % 4
        if pad:
            # Matches FeatureExtractor.padding_value
            input_features = F.pad(input_features, (0, 0, 0, pad), value=1000.0)
            if mask is not None:
                mask = F.pad(mask, (0, pad), value=False)

        logits = self._exported_model(input_features)

        if mask is not None:
            output_mask = self.compute_mask(mask)
            return {"logits": logits, "mask": output_mask}

        if pad:
            implicit_mask = torch.ones(
                batch, input_steps, dtype=torch.bool, device=input_features.device
            )
            valid_len = self.compute_mask(implicit_mask).shape[1]
            logits = logits[:, :valid_len, :]

        return logits

    def to(self, device, *args, **kwargs):
        """Override to() to also move the exported model."""
        super().to(device, *args, **kwargs)
        self._load_exported_model(device)
        return self

    def cuda(self, device=None):
        """Override cuda() to also move the exported model."""
        super().cuda(device)
        self._load_exported_model("cuda")
        return self

    def cpu(self):
        """Override cpu() to also move the exported model."""
        super().cpu()
        self._load_exported_model("cpu")
        return self