import atexit
import json
import os
import sys
from typing import Any, Dict, List, Optional, Union

import numpy as np

# Add models directory to path for imports
_script_dir = os.path.dirname(os.path.abspath(__file__))
_dvlm_dir = os.path.dirname(os.path.dirname(_script_dir))
_models_dir = os.path.join(_dvlm_dir, "models")
if _models_dir not in sys.path:
    sys.path.insert(0, _models_dir)

from nemotron_ministral_mdm_model import NemotronMinistralMDMModel


def _compute_nfe_stats(records: List[Dict[str, Any]]) -> Dict[str, Any]:
    if not records:
        return {}

    tokens = np.array([max(0, int(item.get("generated_tokens", 0))) for item in records])
    nfe = np.array([max(0, int(item.get("nfe", 0))) for item in records])

    with np.errstate(divide="ignore", invalid="ignore"):
        ratio = np.where(nfe > 0, tokens / nfe, 0.0)

    return {
        "total_samples": int(len(records)),
        "total_tokens": int(tokens.sum()),
        "total_nfe": int(nfe.sum()),
        "tokens_per_nfe": float(tokens.sum() / nfe.sum()) if nfe.sum() > 0 else 0.0,
        "tokens_per_nfe_distribution": {
            "mean": float(ratio.mean()) if ratio.size > 0 else 0.0,
            "std": float(ratio.std()) if ratio.size > 0 else 0.0,
            "min": float(ratio.min()) if ratio.size > 0 else 0.0,
            "max": float(ratio.max()) if ratio.size > 0 else 0.0,
            "p50": float(np.percentile(ratio, 50)) if ratio.size > 0 else 0.0,
            "p90": float(np.percentile(ratio, 90)) if ratio.size > 0 else 0.0,
            "p99": float(np.percentile(ratio, 99)) if ratio.size > 0 else 0.0,
        },
    }


class NemotronMinistralChatMDM:
    """VLMEvalKit wrapper for Nemotron-Ministral MDM checkpoints."""

    def __init__(
        self,
        model_path: str,
        processor_path: Optional[str] = None,
        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,
        track_nfe_stats: bool = True,
        **kwargs,
    ):
        del processor_path  # Kept for compatibility with existing config schema.

        self.track_nfe_stats = track_nfe_stats
        self._stats_output_path = kwargs.pop(
            "tokens_per_step_output",
            os.environ.get("TOKENS_PER_STEP_OUTPUT", ""),
        )
        self._stats_dump_every = int(os.environ.get("TOKENS_PER_STEP_DUMP_EVERY", "1"))
        self._records: List[Dict[str, Any]] = []

        self.model = NemotronMinistralMDMModel(
            model_name=model_path,
            device=device,
            block_size=block_size,
            max_tokens=max_tokens,
            steps=steps,
            threshold=threshold,
            shift_logits=shift_logits,
            causal_context=causal_context,
            torch_dtype=kwargs.pop("torch_dtype", "bfloat16"),
            dlm_paradigm=kwargs.pop("dlm_paradigm", "bidirectional"),
        )

        if self.track_nfe_stats:
            atexit.register(self._print_final_stats)

    def _extract_prompt_image(
        self, message: Union[str, Dict[str, Any], List[Any]]
    ) -> Dict[str, Optional[str]]:
        prompt_parts: List[str] = []
        image: Optional[str] = None

        def _read_text(item: Dict[str, Any]) -> Optional[str]:
            value = item.get("value", item.get("text"))
            if value is None:
                return None
            value = str(value).strip()
            return value or None

        def _read_image(item: Dict[str, Any]) -> Optional[str]:
            for key in ("value", "image", "path", "url"):
                value = item.get(key)
                if value is not None and str(value).strip():
                    return str(value)
            return None

        if isinstance(message, str):
            prompt_parts.append(message)
        elif isinstance(message, dict):
            content = message.get("content", "")
            if isinstance(content, str):
                prompt_parts.append(content)
            elif isinstance(content, list):
                for block in content:
                    if not isinstance(block, dict):
                        continue
                    block_type = block.get("type")
                    if block_type == "text":
                        text = _read_text(block)
                        if text:
                            prompt_parts.append(text)
                    elif block_type in {"image", "image_url"}:
                        candidate = _read_image(block)
                        if candidate:
                            image = candidate
        elif isinstance(message, list):
            for item in message:
                if not isinstance(item, dict):
                    continue
                item_type = item.get("type")
                if item_type == "text":
                    text = _read_text(item)
                    if text:
                        prompt_parts.append(text)
                elif item_type in {"image", "image_url"}:
                    candidate = _read_image(item)
                    if candidate:
                        image = candidate

        prompt = " ".join(part.strip() for part in prompt_parts if part and part.strip()).strip()
        return {"prompt": prompt or None, "image": image}

    def generate(
        self,
        message=None,
        dataset=None,
        prompt: Optional[Union[str, List[str]]] = None,
        image: Optional[Union[str, List[str]]] = None,
        **kwargs,
    ) -> Union[str, List[str]]:
        del dataset
        del kwargs

        if message is not None:
            parsed = self._extract_prompt_image(message)
            prompt = parsed["prompt"] if prompt is None else prompt
            image = parsed["image"] if image is None else image

        if prompt is None:
            raise ValueError("No prompt found in message or prompt parameter")

        if isinstance(prompt, list):
            images = (
                image
                if isinstance(image, list)
                else [image] * len(prompt)
                if image is not None
                else [None] * len(prompt)
            )
            return self.model.generate_batch(prompts=prompt, images=images)

        single_image: Optional[str]
        if isinstance(image, list):
            single_image = image[0] if image else None
        else:
            single_image = image

        if self.track_nfe_stats:
            text, stats = self.model.generate(
                prompt=prompt,
                image=single_image,
                return_stats=True,
            )
            if stats:
                self._records.append(stats)
                self._maybe_dump_stats()
            return text

        return self.model.generate(
            prompt=prompt,
            image=single_image,
            return_stats=False,
        )

    def _get_stats(self) -> Dict[str, Any]:
        return _compute_nfe_stats(self._records)

    def _print_final_stats(self):
        if not self._records:
            return
        stats = self._get_stats()
        print("\n" + "=" * 60)
        print("Nemotron-Ministral Decoding Stats")
        print("=" * 60)
        print(f"Total samples: {stats.get('total_samples', 0)}")
        print(f"Total tokens:  {stats.get('total_tokens', 0)}")
        print(f"Total NFE:     {stats.get('total_nfe', 0)}")
        print(f"Tokens / NFE:  {stats.get('tokens_per_nfe', 0):.4f}")
        dist = stats.get("tokens_per_nfe_distribution", {})
        print(
            f"Distribution:  mean={dist.get('mean', 0):.4f}, std={dist.get('std', 0):.4f}, "
            f"p50={dist.get('p50', 0):.4f}, p90={dist.get('p90', 0):.4f}, p99={dist.get('p99', 0):.4f}"
        )
        print("=" * 60)
        self._write_stats_file(stats)

    def _maybe_dump_stats(self):
        if not self._stats_output_path:
            return
        if self._stats_dump_every <= 0:
            return
        if len(self._records) % self._stats_dump_every != 0:
            return
        self._write_stats_file(self._get_stats())

    def _write_stats_file(self, stats: Dict[str, Any]):
        if not self._stats_output_path:
            return
        output_dir = os.path.dirname(self._stats_output_path)
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)
        with open(self._stats_output_path, "w", encoding="utf-8") as file:
            file.write(json.dumps(stats, ensure_ascii=False, indent=2))

    def set_dump_image(self, dump_image: bool):
        self.dump_image = dump_image


# Register with VLMEvalKit if possible
try:
    import vlmeval.api as api

    if not hasattr(api, "NemotronMinistralChatMDM"):
        api.NemotronMinistralChatMDM = NemotronMinistralChatMDM
except ImportError:
    pass

try:
    import vlmeval.vlm as vlm

    if not hasattr(vlm, "NemotronMinistralChatMDM"):
        vlm.NemotronMinistralChatMDM = NemotronMinistralChatMDM
except ImportError:
    pass
