# Prompt: 完全按照chatbot_mdm.py的方式构造执行
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 qwen_vl_mdm_model import QwenVLMDMModel


def _compute_tokens_per_step_stats(all_tokens_per_step: List[int]) -> Dict[str, Any]:
    """Compute statistics for tokens decoded per step."""
    if not all_tokens_per_step:
        return {}
    arr = np.array(all_tokens_per_step)
    total_tokens = int(arr.sum())
    total_steps = len(arr)
    
    # Distribution: how many steps decoded 1/2/3+ tokens
    one_tok = int((arr == 1).sum())
    two_tok = int((arr == 2).sum())
    three_plus = int((arr >= 3).sum())
    
    return {
        "mean": float(arr.mean()),
        "std": float(arr.std()),
        "min": int(arr.min()),
        "max": int(arr.max()),
        "p50": float(np.percentile(arr, 50)),
        "p90": float(np.percentile(arr, 90)),
        "p99": float(np.percentile(arr, 99)),
        "total_tokens": total_tokens,
        "total_steps": total_steps,
        "tokens_per_step_ratio": total_tokens / total_steps if total_steps > 0 else 0,
        "distribution": {
            "1_token": one_tok,
            "2_tokens": two_tok,
            "3+_tokens": three_plus,
        },
        "distribution_pct": {
            "1_token": 100.0 * one_tok / total_steps if total_steps > 0 else 0,
            "2_tokens": 100.0 * two_tok / total_steps if total_steps > 0 else 0,
            "3+_tokens": 100.0 * three_plus / total_steps if total_steps > 0 else 0,
        },
    }


class Qwen2VLChatMDM:
    """VLMEvalKit wrapper using QwenVLMDMModel (same as chatbot_mdm_v2.py).
    
    Supports A/B testing between standard and speculative decoding:
    - use_speculative=False (default): Standard block-causal decoding
    - use_speculative=True: Speculative block-causal decoding (faster)
    
    Supports A/B testing between L2R and confidence-based decoding:
    - force_left_to_right=True (default): Sequential L2R decoding
    - force_left_to_right=False: Confidence-based parallel unmasking
    """

    def __init__(
        self,
        model_path: str,
        processor_path: Optional[str] = None,
        device: str = "cuda",
        block_size: Optional[int] = None,
        max_tokens: int = 128,
        mask_token: str = "|<MASK>|",
        use_speculative: bool = False,  # A/B test toggle
        threshold: Optional[float] = None,
        force_left_to_right: bool = True,  # A/B test: L2R vs confidence-based
        track_tokens_per_step: bool = True,  # Track decoding stats (default on)
        **kwargs,
    ):
        self.use_speculative = use_speculative
        self.force_left_to_right = force_left_to_right
        self.track_tokens_per_step = track_tokens_per_step
        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")
        )
        
        # Stats tracking
        self._all_tokens_per_step: List[int] = []
        self._total_samples = 0
        
        self.model = QwenVLMDMModel(
            model_name=model_path,
            processor_name=processor_path or model_path,
            device=device,
            block_size=block_size,  # None = use config.bd_size
            max_tokens=max_tokens,
            mask_token=mask_token,
            use_speculative=use_speculative,
            threshold=threshold,
            force_left_to_right=force_left_to_right,
        )
        
        # Register atexit handler to print stats when process ends
        if self.track_tokens_per_step:
            atexit.register(self._print_final_stats)

    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]]:
        # Handle VLMEvalKit calling style: generate(message=struct, dataset=dataset_name)
        if message is not None:
            # Extract prompt and image from message structure
            # Completely reconstruct according to chatbot_mdm.py input convention:
            # Accept message as list of {'type': ..., 'value': ...} dicts (as in the prompt above)
            # Reconstruct prompt and image in the same way as chatbot_mdm.py

            # Allow passing message as list of dicts (like [{'type': 'image', ...}, {'type': 'text', ...}])
            if isinstance(message, list):
                content = []
                image_path = None
                prompt_parts = []
                for item in message:
                    if isinstance(item, dict):
                        if item.get("type") == "image" and "value" in item:
                            image_path = item["value"]
                        elif item.get("type") == "text" and "value" in item:
                            prompt_parts.append(item["value"])
                if prompt_parts:
                    prompt = " ".join(prompt_parts)
                if image_path:
                    image = image_path
            else:
                raise ValueError(f"Unsupported message type: {type(message)}")

        # Fallback to prompt/image if message extraction failed
        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 else [None] * len(prompt)
            return self.model.generate_batch(prompts=prompt, images=images)
        else:
            img = image if isinstance(image, str) else image[0] if image else None
            
            # Track tokens_per_step if enabled and not using speculative decoding
            if self.track_tokens_per_step and not self.use_speculative:
                result, stats = self.model.generate(prompt=prompt, image=img, return_stats=True)
                if stats and "tokens_per_step" in stats:
                    self._all_tokens_per_step.extend(stats["tokens_per_step"])
                self._total_samples += 1
                self._maybe_dump_stats()
                return result
            else:
                return self.model.generate(prompt=prompt, image=img)
    
    def get_tokens_per_step_stats(self) -> Dict[str, Any]:
        """Get accumulated tokens_per_step statistics."""
        stats = _compute_tokens_per_step_stats(self._all_tokens_per_step)
        stats["total_samples"] = self._total_samples
        return stats
    
    def _print_final_stats(self):
        """Print final statistics (called at process exit)."""
        if not self._all_tokens_per_step:
            return
        
        stats = self.get_tokens_per_step_stats()
        print("\n" + "=" * 60)
        print("MDM Tokens Per Step Statistics")
        print("=" * 60)
        print(f"Total samples: {stats.get('total_samples', 0)}")
        print(f"Total tokens:  {stats.get('total_tokens', 0)}")
        print(f"Total steps:   {stats.get('total_steps', 0)}")
        print(f"Tokens/Step:   mean={stats.get('mean', 0):.3f}, std={stats.get('std', 0):.3f}")
        print(f"               min={stats.get('min', 0)}, max={stats.get('max', 0)}")
        print(f"               p50={stats.get('p50', 0):.2f}, p90={stats.get('p90', 0):.2f}, p99={stats.get('p99', 0):.2f}")
        dist_pct = stats.get('distribution_pct', {})
        print(f"Distribution:  1-tok={dist_pct.get('1_token', 0):.1f}%, "
              f"2-tok={dist_pct.get('2_tokens', 0):.1f}%, "
              f"3+tok={dist_pct.get('3+_tokens', 0):.1f}%")
        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 self._total_samples % self._stats_dump_every != 0:
            return
        stats = self.get_tokens_per_step_stats()
        self._write_stats_file(stats)

    def _write_stats_file(self, stats: Dict[str, Any]):
        if not self._stats_output_path:
            return
        os.makedirs(os.path.dirname(self._stats_output_path), 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):
        """Set whether to dump images (required by VLMEvalKit)."""
        self.dump_image = dump_image


# Register with VLMEvalKit if possible
try:
    import vlmeval.api as api
    if not hasattr(api, 'Qwen2VLChatMDM'):
        api.Qwen2VLChatMDM = Qwen2VLChatMDM
except ImportError:
    pass

try:
    import vlmeval.vlm as vlm
    if not hasattr(vlm, 'Qwen2VLChatMDM'):
        vlm.Qwen2VLChatMDM = Qwen2VLChatMDM
except ImportError:
    pass

