"""Core Nemotron-Ministral MDM model wrapper for evaluation/inference."""

from __future__ import annotations

import os
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union

import torch
from transformers import AutoConfig, AutoModel, AutoTokenizer


def _resolve_path(path: str) -> str:
    if os.path.exists(path):
        return os.path.abspath(path)
    return path


def _resolve_torch_dtype(dtype: Optional[Union[str, torch.dtype]]) -> Optional[torch.dtype]:
    if dtype is None:
        return None
    if isinstance(dtype, torch.dtype):
        return dtype
    dtype_str = str(dtype).lower().strip()
    if dtype_str in {"", "auto", "none"}:
        return None
    mapping = {
        "bfloat16": torch.bfloat16,
        "bf16": torch.bfloat16,
        "float16": torch.float16,
        "fp16": torch.float16,
        "float32": torch.float32,
        "fp32": torch.float32,
    }
    if dtype_str not in mapping:
        raise ValueError(f"Unsupported torch dtype: {dtype}")
    return mapping[dtype_str]


class NemotronMinistralMDMModel:
    """Nemotron-Ministral wrapper using tokenizer.process_messages + model.generate."""

    def __init__(
        self,
        model_name: str,
        device: str = "cuda",
        block_size: Optional[int] = None,
        max_tokens: int = 512,
        steps: Optional[int] = None,
        threshold: float = 0.9,
        shift_logits: bool = False,
        causal_context: bool = True,
        add_generation_prompt: bool = True,
        torch_dtype: Optional[Union[str, torch.dtype]] = "bfloat16",
        dlm_paradigm: Optional[str] = "bidirectional",
    ):
        self.model_name = _resolve_path(model_name)
        self.device = self._resolve_device(device)
        self.add_generation_prompt = add_generation_prompt
        self.default_threshold = float(threshold)
        self.default_shift_logits = bool(shift_logits)
        self.default_causal_context = bool(causal_context)
        self.default_max_tokens = int(max_tokens)
        self.default_steps = int(steps) if steps is not None else None

        resolved_dtype = _resolve_torch_dtype(torch_dtype)

        self.tokenizer = AutoTokenizer.from_pretrained(
            self.model_name, trust_remote_code=True
        )
        self.config = AutoConfig.from_pretrained(self.model_name, trust_remote_code=True)
        if dlm_paradigm is not None:
            self.config.dlm_paradigm = dlm_paradigm

        model_kwargs: Dict[str, Any] = {"trust_remote_code": True, "config": self.config}
        if resolved_dtype is not None:
            model_kwargs["torch_dtype"] = resolved_dtype

        if self.device.startswith("cuda") and torch.cuda.is_available():
            torch.cuda.set_device(0)

        self.model = AutoModel.from_pretrained(self.model_name, **model_kwargs).eval()
        self.model = self.model.to(self.device)
        if resolved_dtype is not None:
            self.model = self.model.to(resolved_dtype)

        inferred_block_size = self._infer_block_size()
        self.default_block_size = int(block_size) if block_size else inferred_block_size

    @staticmethod
    def _resolve_device(device: str) -> str:
        if device.startswith("cuda"):
            return "cuda:0"
        return device

    def _infer_block_size(self) -> int:
        for attr_name in ("block_size", "bd_size"):
            value = getattr(self.config, attr_name, None)
            if isinstance(value, int) and value > 0:
                return value
        return 32

    @staticmethod
    def _normalize_lengths(
        max_tokens: int, steps: int, block_size: int
    ) -> Tuple[int, int, int]:
        max_tokens = max(1, int(max_tokens))
        steps = max(1, int(steps))
        block_size = max(1, int(block_size))

        if max_tokens % block_size != 0:
            adjusted = ((max_tokens + block_size - 1) // block_size) * block_size
            print(
                f"[NemotronMinistralMDMModel] Adjust max_tokens {max_tokens} -> {adjusted} "
                f"to match block_size={block_size}"
            )
            max_tokens = adjusted

        num_blocks = max_tokens // block_size
        if steps % num_blocks != 0:
            adjusted_steps = ((steps + num_blocks - 1) // num_blocks) * num_blocks
            print(
                f"[NemotronMinistralMDMModel] Adjust steps {steps} -> {adjusted_steps} "
                f"to satisfy steps % num_blocks == 0 (num_blocks={num_blocks})"
            )
            steps = adjusted_steps

        return max_tokens, steps, block_size

    def _eos_token_ids(self) -> List[int]:
        eos: Optional[Union[int, Sequence[int]]] = None
        if getattr(self.model, "generation_config", None) is not None:
            eos = getattr(self.model.generation_config, "eos_token_id", None)
        if eos is None:
            eos = getattr(self.tokenizer, "eos_token_id", None)
        if eos is None:
            return []
        if isinstance(eos, (list, tuple, set)):
            return [int(token_id) for token_id in eos if token_id is not None]
        return [int(eos)]

    @staticmethod
    def _trim_to_eos(sequence: torch.Tensor, eos_ids: List[int]) -> torch.Tensor:
        if not eos_ids or sequence.numel() == 0:
            return sequence
        eos_set = set(eos_ids)
        for idx, token_id in enumerate(sequence.tolist()):
            if token_id in eos_set:
                return sequence[:idx]
        return sequence

    def build_messages(self, image: Optional[str], prompt: str) -> List[Dict[str, Any]]:
        content: List[Dict[str, Any]] = []
        if image:
            content.append({"type": "image", "image": image})
        content.append({"type": "text", "text": prompt})
        return [{"role": "user", "content": content}]

    def _prepare_inputs(
        self, prompt: str, image: Optional[str]
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
        messages = self.build_messages(image=image, prompt=prompt)
        processed = self.tokenizer.process_messages(
            messages, add_generation_prompt=self.add_generation_prompt, enable_thinking=False
        )

        prompt_ids = processed["input_ids"]
        if not torch.is_tensor(prompt_ids):
            prompt_ids = torch.as_tensor(prompt_ids)
        prompt_ids = prompt_ids.to(device=self.device)

        pixel_values: Optional[torch.Tensor] = None
        image_sizes: Optional[torch.Tensor] = None

        if "pixel_values" in processed:
            pixel_values = processed["pixel_values"]
            if not torch.is_tensor(pixel_values):
                pixel_values = torch.as_tensor(pixel_values)
            target_dtype = self.model.dtype if self.model.dtype.is_floating_point else torch.bfloat16
            pixel_values = pixel_values.to(device=self.device, dtype=target_dtype)

        if "image_sizes" in processed:
            image_sizes = torch.as_tensor(
                processed["image_sizes"], dtype=torch.long, device=self.device
            )

        return prompt_ids, pixel_values, image_sizes

    @torch.no_grad()
    def generate(
        self,
        prompt: str,
        image: Optional[str] = None,
        max_tokens: Optional[int] = None,
        steps: Optional[int] = None,
        block_size: Optional[int] = None,
        threshold: Optional[float] = None,
        shift_logits: Optional[bool] = None,
        causal_context: Optional[bool] = None,
        return_stats: bool = False,
    ) -> Union[str, Tuple[str, Dict[str, Any]]]:
        max_tokens_final = int(max_tokens) if max_tokens is not None else self.default_max_tokens
        steps_final = int(steps) if steps is not None else (
            self.default_steps if self.default_steps is not None else max_tokens_final
        )
        block_size_final = int(block_size) if block_size is not None else self.default_block_size
        threshold_final = float(threshold) if threshold is not None else self.default_threshold
        shift_logits_final = (
            bool(shift_logits) if shift_logits is not None else self.default_shift_logits
        )
        causal_context_final = (
            bool(causal_context) if causal_context is not None else self.default_causal_context
        )

        max_tokens_final, steps_final, block_size_final = self._normalize_lengths(
            max_tokens=max_tokens_final,
            steps=steps_final,
            block_size=block_size_final,
        )

        prompt_ids, pixel_values, image_sizes = self._prepare_inputs(prompt=prompt, image=image)
        out_ids, nfe = self.model.generate(
            prompt_ids,
            max_new_tokens=max_tokens_final,
            steps=steps_final,
            block_length=block_size_final,
            shift_logits=shift_logits_final,
            causal_context=causal_context_final,
            threshold=threshold_final,
            pixel_values=pixel_values,
            image_sizes=image_sizes,
        )

        response_ids = out_ids[:, prompt_ids.shape[1]:]
        eos_ids = self._eos_token_ids()
        first_response = self._trim_to_eos(response_ids[0], eos_ids)
        output_text = self.tokenizer.decode(
            first_response.tolist(),
            skip_special_tokens=True,
        ).strip()

        if return_stats:
            generated_tokens = int(first_response.numel())
            nfe_int = int(nfe)
            stats = {
                "prompt_tokens": int(prompt_ids.shape[1]),
                "generated_tokens": generated_tokens,
                "nfe": nfe_int,
                "tokens_per_nfe": (
                    float(generated_tokens) / float(nfe_int) if nfe_int > 0 else 0.0
                ),
                "block_size": block_size_final,
                "steps": steps_final,
                "max_tokens": max_tokens_final,
                "threshold": threshold_final,
            }
            return output_text, stats

        return output_text

    def generate_batch(
        self,
        prompts: List[str],
        images: Optional[List[Optional[str]]] = None,
        max_tokens: Optional[int] = None,
        steps: Optional[int] = None,
        block_size: Optional[int] = None,
        threshold: Optional[float] = None,
    ) -> List[str]:
        if images is None:
            images = [None] * len(prompts)
        outputs: List[str] = []
        for prompt, image in zip(prompts, images):
            outputs.append(
                self.generate(
                    prompt=prompt,
                    image=image,
                    max_tokens=max_tokens,
                    steps=steps,
                    block_size=block_size,
                    threshold=threshold,
                    return_stats=False,
                )
            )
        return outputs
