# Prompt: 弄成input parameters
# Core model class for Qwen-VL with MDM sampling, extracted from chatbot_mdm.py
import os
import types
from typing import Any, Dict, List, Optional, Tuple

import torch
from transformers import AutoModel, AutoProcessor, AutoTokenizer

from generation_functions import Fast_dLLM_QwenForCausalLM

try:
    from qwen_vl_utils import process_vision_info
except ImportError:
    # Fallback if qwen_vl_utils is not available
    def process_vision_info(messages):
        """Fallback implementation if qwen_vl_utils is not available."""
        image_inputs = []
        video_inputs = []
        for msg in messages:
            if isinstance(msg.get("content"), list):
                for item in msg["content"]:
                    if item.get("type") == "image":
                        image_inputs.append(item.get("image"))
                    elif item.get("type") == "video":
                        video_inputs.append(item.get("video"))
        return image_inputs if image_inputs else None, video_inputs if video_inputs else None


class QwenVLMDMModel:
    """
    Core model class that encapsulates the model loading and MDM sampling logic
    from chatbot_mdm.py. This class can be reused by VLMEvalKit wrapper.
    
    Supports two decoding modes for A/B testing:
    - Standard block-causal decoding (use_speculative=False, default)
    - Speculative block-causal decoding (use_speculative=True)
    
    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_name: str,
        processor_name: Optional[str] = None,
        device: str = "cuda",
        block_size: Optional[int] = None,
        max_tokens: int = 128,
        mask_token: str = "|<MASK>|",
        threshold: Optional[float] = None,
        use_speculative: bool = False,  # A/B test: False=standard, True=speculative
        force_left_to_right: bool = True,  # A/B test: True=L2R sequential, False=confidence-based
    ):
        """
        Initialize the Qwen-VL MDM model.

        Args:
            model_name: Model identifier or local path for AutoModel/AutoTokenizer
            processor_name: Processor identifier for AutoProcessor (defaults to model_name)
            device: Device to move processor outputs to (e.g., 'cuda' or 'cpu')
            block_size: Block size passed to model.mdm_sample (None=use config.bd_size)
            max_tokens: Maximum number of tokens to generate
            mask_token: Mask token string used to compute mask_id
        """
        # Convert relative paths to absolute paths for local checkpoints
        if os.path.exists(model_name):
            model_name = os.path.abspath(model_name)
        self.model_name = model_name
        
        # For processor, also convert to absolute if it's a local path
        if processor_name and os.path.exists(processor_name):
            processor_name = os.path.abspath(processor_name)
        self.processor_name = processor_name or model_name
        self._block_size_override = block_size  # Store override value
        # Resolve device based on torchrun LOCAL_RANK if not explicitly specified
        resolved_device = device
        # In torchrun with CUDA_VISIBLE_DEVICES set per process,
        # each process only sees one GPU as cuda:0
        if device.startswith("cuda"):
            resolved_device = "cuda:0"

        self.device = resolved_device
        self.max_tokens = max_tokens
        self.mask_token = mask_token
        self.threshold = threshold
        if threshold is not None:
            self.threshold = threshold
        else:
            self.threshold = 1
        
        # A/B test flag: select between standard and speculative decoding
        self.use_speculative = use_speculative
        # A/B test flag: select between L2R sequential and confidence-based decoding
        self.force_left_to_right = force_left_to_right
        # Load model
        # In multi-GPU torchrun, CUDA_VISIBLE_DEVICES is set per process
        # Each process only sees one GPU as cuda:0
        import torch
        
        # Ensure we're using the first (and only) visible GPU
        if torch.cuda.is_available():
            torch.cuda.set_device(0)
        
        from lmflow.models.qwen2_5_vl_mdm import Fast_dLLM_Qwen2_5_VLForConditionalGeneration
        self.model = Fast_dLLM_Qwen2_5_VLForConditionalGeneration.from_pretrained(
            self.model_name,
            torch_dtype="auto",
            trust_remote_code=True,
            device_map="cuda:0",
        )
        # self.model = AutoModel.from_pretrained(
        #     self.model_name,
        #     torch_dtype=torch.bfloat16,
        #     device_map="cuda:0",  # Explicitly use cuda:0 (the only visible GPU)
        #     trust_remote_code=True,
        # )
        # Use override block_size if provided, otherwise use config.bd_size
        # For AR models (no bd_size in config), default to block_size=1
        if self._block_size_override is not None:
            self.block_size = self._block_size_override
        elif hasattr(self.model.config, 'bd_size') and self.model.config.bd_size is not None:
            self.block_size = self.model.config.bd_size
        else:
            # AR model or MDM model without bd_size - default to 1 (standard AR decoding)
            self.block_size = 1
            print(f"[QwenVLMDMModel] Warning: No bd_size in config, using block_size=1 (AR mode)")

        # Bind mdm_sample methods for A/B testing
        # Standard block-causal decoding (A)
        self.model.mdm_sample = types.MethodType(
            Fast_dLLM_QwenForCausalLM.block_causal_sample_cache, self.model
        )
        # Speculative block-causal decoding (B)
        self.model.speculative_sample = types.MethodType(
            Fast_dLLM_QwenForCausalLM.speculative_block_causal_sample_cache, self.model
        )

        # Handle rope_deltas (from chatbot_mdm.py lines 79-93)
        # language_model = getattr(self.model, "model", None)
        # if language_model is not None and hasattr(language_model, "language_model"):
        #     lm = language_model.language_model
        #     if not hasattr(lm, "_rope_safe_forward"):
        #         original_forward = lm.forward

        #         def forward_with_rope(self, *method_args, **method_kwargs):
        #             outputs = original_forward(*method_args, **method_kwargs)
        #             if not hasattr(outputs, "rope_deltas"):
        #                 rope_deltas = getattr(language_model, "rope_deltas", None)
        #                 setattr(outputs, "rope_deltas", rope_deltas)
        #             return outputs

        #         lm.forward = types.MethodType(forward_with_rope, lm)
        #         lm._rope_safe_forward = original_forward

        # Load processor and tokenizer
        self.processor = AutoProcessor.from_pretrained(
            self.processor_name, use_fast=False
        )
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.processor.tokenizer = self.tokenizer

        # Compute mask_id
        self.mask_id = self.tokenizer.encode(self.mask_token)[0]

    def build_messages(
        self, image: Optional[str], prompt: str
    ) -> List[Dict[str, Any]]:
        """
        Build messages in the format expected by the processor.

        Args:
            image: Image URL or path (None to skip image input)
            prompt: Text prompt

        Returns:
            List of message dictionaries
        """
        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 generate(
        self,
        prompt: str,
        image: Optional[str] = None,
        max_tokens: Optional[int] = None,
        block_size: Optional[int] = None,
        threshold: Optional[float] = None,
        return_stats: bool = False,
        debug: bool = False,
        output_file: Optional[str] = None,
    ) -> str:
        """
        Generate text using MDM sampling.

        Args:
            prompt: Text prompt
            image: Optional image URL or path
            max_tokens: Override default max_tokens
            block_size: Override default block_size
            threshold: Override default threshold
            return_stats: If True, return (text, stats_dict) with tokens_per_step info
            debug: If True, write step-by-step decoding trace to output_file
            output_file: Path for debug trace (default: decoded_output_v2.txt)

        Returns:
            Generated text, or (text, stats) if return_stats=True
        """
        # import pdb; pdb.set_trace()

        block_size = block_size or self.block_size
        max_tokens = max_tokens or self.max_tokens
        threshold = threshold or self.threshold
        # import pdb; pdb.set_trace()
        # Build messages
        messages = self.build_messages(image, prompt)

        # Apply chat template
        text = self.processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )

        # Process vision info
        image_inputs, video_inputs = process_vision_info(messages)

        # Process inputs
        inputs = self.processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        inputs = inputs.to(self.device)

        # Prepare mdm_kwargs
        mdm_kwargs = {
            "input_ids": inputs.input_ids,
            "tokenizer": self.processor.tokenizer,
            "block_size": block_size,
            "max_tokens": max_tokens,
            "mask_id": self.mask_id,
            "return_stats": return_stats,
            "threshold": threshold,
            "force_left_to_right": self.force_left_to_right,
            "sub_block_size": 4,
            "debug": debug,
            "output_file": output_file or "decoded_output_v2.txt",
        }
        if hasattr(inputs, "pixel_values"):
            mdm_kwargs["pixel_values"] = inputs.pixel_values
        if hasattr(inputs, "image_grid_thw"):
            mdm_kwargs["image_grid_thw"] = inputs.image_grid_thw

        # Generate using the selected decoding method (A/B test)
        if self.use_speculative:
            # Speculative decoding doesn't support return_stats or force_left_to_right
            mdm_kwargs.pop('return_stats', None)
            mdm_kwargs.pop('force_left_to_right', None)
            # threshold=0: accept all predictions, let AR verification handle correctness
            mdm_kwargs['threshold'] = 0
            result = self.model.speculative_sample(**mdm_kwargs)
            stats = None
        else:
            result = self.model.mdm_sample(**mdm_kwargs)
        
        if return_stats:
            generated_ids, stats = result
        else:
            generated_ids = result
            stats = None

        # Trim generated tokens (remove input tokens)
        generated_ids_trimmed = [
            out_ids[len(in_ids) :]
            for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
        ]

        # Decode output
        output_text = self.processor.batch_decode(
            generated_ids_trimmed,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        )

        text_output = output_text[0] if output_text else ""
        
        if return_stats:
            return text_output, stats
        return text_output

    def generate_batch(
        self,
        prompts: List[str],
        images: Optional[List[Optional[str]]] = None,
        max_tokens: Optional[int] = None,
        block_size: Optional[int] = None,
        threshold: Optional[float] = None,
    ) -> List[str]:
        """
        Generate text for multiple prompts (batch processing).

        Args:
            prompts: List of text prompts
            images: Optional list of image URLs/paths (same length as prompts)
            max_tokens: Override default max_tokens
            block_size: Override default block_size

        Returns:
            List of generated texts
        """
        if images is None:
            images = [None] * len(prompts)
        threshold = threshold or self.threshold
        results = []
        for prompt, image in zip(prompts, images):
            result = self.generate(
                prompt=prompt,
                image=image,
                max_tokens=max_tokens,
                block_size=block_size,
                threshold=threshold,
            )
            results.append(result)

        return results

