#!/usr/bin/env python3
"""
Unit tests for VLM/LLM alignment: verifying that the VLM training script
and DataCollatorForQwenVL behave correctly after alignment cleanup.

Tests cover:
1. Shell script parameter alignment (VLM vs LLM)
2. DataCollatorForQwenVL correctness
3. Model forward path dispatch (LLM vs VLM)
4. finetune_qwenvl.py worker cap logic
5. MDM config propagation

#cpu - all tests run on CPU only
"""

import os
import re
import subprocess
import sys
import pytest
import torch
from unittest.mock import MagicMock, patch
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent


# ============================================================
# 1. Shell script alignment tests
# ============================================================

class TestShellScriptAlignment:
    """Verify VLM training script is aligned with LLM script."""

    VLM_SCRIPT = PROJECT_ROOT / "chengyue_script/dvlm/training/multinode/mdm/j63_fastdllmv2/train_8node_bsz256_anneal_bd32_block_causal_lr_1e-6_unfreeze_vision.sh"
    LLM_SCRIPT = PROJECT_ROOT / "chengyue_script/dvlm/training/single_node/run_shizhe_no_reasoning_mdm_group_pad_complementary_seq_length_2k_mdm_block_causal.sh"

    @pytest.fixture(autouse=True)
    def load_scripts(self):
        """#cpu"""
        self.vlm_content = self.VLM_SCRIPT.read_text()
        if self.LLM_SCRIPT.exists():
            self.llm_content = self.LLM_SCRIPT.read_text()
        else:
            self.llm_content = ""

    def test_vlm_script_exists(self):
        """#cpu"""
        assert self.VLM_SCRIPT.exists(), f"VLM script missing: {self.VLM_SCRIPT}"

    def test_workers_aligned_to_8(self):
        """#cpu - WORKERS should default to 8, not 1."""
        match = re.search(r'WORKERS="\$\{WORKERS:-(\d+)\}"', self.vlm_content)
        assert match is not None, "WORKERS default not found"
        assert match.group(1) == "8", f"WORKERS default should be 8, got {match.group(1)}"

    def test_no_dead_env_var_mdm_vl_keep_vision(self):
        """#cpu - MDM_VL_KEEP_VISION is a dead variable, should not appear."""
        assert "MDM_VL_KEEP_VISION" not in self.vlm_content, \
            "MDM_VL_KEEP_VISION is dead code (never read by Python), should be removed"

    def test_no_report_to_none(self):
        """#cpu - --report_to none was VLM-specific bloat."""
        assert "--report_to none" not in self.vlm_content, \
            "--report_to none should be removed to align with LLM"

    def test_has_validation_split_percentage(self):
        """#cpu - Should have --validation_split_percentage 0 like LLM."""
        assert "--validation_split_percentage 0" in self.vlm_content or \
               "--validation_split_percentage" in self.vlm_content, \
            "Missing --validation_split_percentage (LLM has it)"

    def test_uses_finetune_qwenvl(self):
        """#cpu - VLM must use finetune_qwenvl.py, not finetune.py."""
        assert "finetune_qwenvl.py" in self.vlm_content, \
            "VLM script must use finetune_qwenvl.py for image processing"

    def test_fsdp_shard_grad_op(self):
        """#cpu - Both scripts should use shard_grad_op FSDP."""
        assert '"shard_grad_op"' in self.vlm_content, "Missing FSDP shard_grad_op"

    def test_gradient_checkpointing_enabled(self):
        """#cpu - Activation checkpointing should be enabled by default."""
        assert "ACTIVATION_CHECKPOINTING" in self.vlm_content
        match = re.search(r'ACTIVATION_CHECKPOINTING="\$\{ACTIVATION_CHECKPOINTING:-(\d+)\}"', self.vlm_content)
        assert match is not None
        assert match.group(1) == "1", "ACTIVATION_CHECKPOINTING should default to 1"

    def test_no_skip_adapt(self):
        """#cpu - SKIP_ADAPT was unused, should be removed."""
        assert "SKIP_ADAPT" not in self.vlm_content, \
            "SKIP_ADAPT is unused and should be removed"

    def test_no_pixel_sanity_check(self):
        """#cpu - MIN_PIXELS/MAX_PIXELS sanity checks were bloat."""
        assert "MIN_PIXELS % 196" not in self.vlm_content, \
            "Pixel sanity check is unnecessary bloat"

    def test_mdm_params_present(self):
        """#cpu - Core MDM params must be present."""
        for param in ["--mdm", "--bd_size", "--complementary_mask",
                       "--always_mask_im_end", "--use_block_causal_mask",
                       "--learn_padding", "--block_size"]:
            assert param in self.vlm_content, f"Missing MDM param: {param}"

    def test_no_max_length_param(self):
        """#cpu - --max_length was removed to align with LLM."""
        # It's ok if MAX_LENGTH variable is defined but NOT passed to torchrun
        # Check that --max_length is not in the torchrun command
        # Find the srun/torchrun section
        torchrun_section = self.vlm_content[self.vlm_content.find("torchrun"):]
        assert "--max_length" not in torchrun_section, \
            "--max_length should be removed from torchrun args to align with LLM"

    def test_signal_handler_present(self):
        """#cpu - Signal handler for SLURM requeue must be present."""
        assert "handle_timeout" in self.vlm_content
        assert "trap" in self.vlm_content
        assert "USR1" in self.vlm_content

    def test_auto_resume_logic(self):
        """#cpu - Auto-resume checkpoint detection must be present."""
        assert "latest_checkpoint" in self.vlm_content
        assert "resume_from_checkpoint" in self.vlm_content


# ============================================================
# 2. DataCollatorForQwenVL tests
# ============================================================

class TestDataCollatorForQwenVL:
    """Test the data collator that handles vision+text preprocessing."""

    @pytest.fixture
    def mock_processor(self):
        """#cpu - Create a mock Qwen2.5-VL processor."""
        proc = MagicMock()
        # apply_chat_template returns raw text
        proc.apply_chat_template.return_value = (
            "<|im_start|>user\nHello.<|im_end|>\n<|im_start|>assistant\nHi!<|im_end|>\n"
        )
        # processor() returns dict with input_ids and attention_mask
        mock_output = {
            "input_ids": torch.tensor([[151644, 77091, 198, 100, 200, 300, 151645, 198]]),
            "attention_mask": torch.ones(1, 8, dtype=torch.long),
        }
        proc.return_value = mock_output
        return proc

    @pytest.fixture
    def mock_tokenizer(self):
        """#cpu"""
        tok = MagicMock()
        tok.encode.return_value = [151645]  # <|im_end|>
        tok.convert_tokens_to_ids.return_value = 151665  # |<MASK>|
        tok.pad_token_id = 151643
        return tok

    def test_collator_text_only_basic(self, mock_processor, mock_tokenizer):
        """#cpu - Text-only instance should produce valid batch."""
        from lmflow.datasets.multi_modal_dataset import DataCollatorForQwenVL

        collator = DataCollatorForQwenVL(
            processor=mock_processor,
            tokenizer=mock_tokenizer,
            bd_size=None,
            max_length=None,
        )
        instances = [
            {"messages": [
                {"role": "user", "content": [{"type": "text", "text": "Hello."}]},
                {"role": "assistant", "content": [{"type": "text", "text": "Hi!"}]},
            ]}
        ]
        batch = collator(instances)

        assert "input_ids" in batch
        assert "labels" in batch
        assert "attention_mask" in batch
        assert batch["input_ids"].ndim == 2
        assert batch["labels"].ndim == 2

    def test_collator_empty_messages_fallback(self, mock_processor, mock_tokenizer):
        """#cpu - Empty messages should trigger dummy fallback."""
        from lmflow.datasets.multi_modal_dataset import DataCollatorForQwenVL

        collator = DataCollatorForQwenVL(
            processor=mock_processor,
            tokenizer=mock_tokenizer,
            bd_size=None,
            max_length=None,
        )
        instances = [{"messages": []}]
        batch = collator(instances)

        assert "input_ids" in batch
        assert batch["input_ids"].shape[0] == 1  # dummy batch of size 1

    def test_collator_bd_size_padding(self, mock_processor, mock_tokenizer):
        """#cpu - With bd_size set, output seq length must be divisible by bd_size."""
        from lmflow.datasets.multi_modal_dataset import DataCollatorForQwenVL

        bd_size = 4
        collator = DataCollatorForQwenVL(
            processor=mock_processor,
            tokenizer=mock_tokenizer,
            bd_size=bd_size,
            max_length=None,
        )
        instances = [
            {"messages": [
                {"role": "user", "content": [{"type": "text", "text": "Hello."}]},
                {"role": "assistant", "content": [{"type": "text", "text": "Hi!"}]},
            ]}
        ]
        batch = collator(instances)

        seq_len = batch["input_ids"].shape[1]
        assert seq_len % bd_size == 0, f"Seq length {seq_len} not divisible by bd_size {bd_size}"

    def test_collator_labels_mask_padding(self, mock_processor, mock_tokenizer):
        """#cpu - Padding positions in labels should be -100."""
        from lmflow.datasets.multi_modal_dataset import DataCollatorForQwenVL

        collator = DataCollatorForQwenVL(
            processor=mock_processor,
            tokenizer=mock_tokenizer,
            bd_size=32,
            max_length=None,
        )
        instances = [
            {"messages": [
                {"role": "user", "content": [{"type": "text", "text": "Hello."}]},
                {"role": "assistant", "content": [{"type": "text", "text": "Hi!"}]},
            ]}
        ]
        batch = collator(instances)

        # The mask_id padding positions should have labels = -100
        mask_id = 151665
        mask_positions = (batch["input_ids"][0] == mask_id)
        if mask_positions.any():
            assert (batch["labels"][0][mask_positions] == -100).all(), \
                "Padding positions (mask_id) should have labels = -100"

    def test_collator_batch_output_shapes(self, mock_processor, mock_tokenizer):
        """#cpu - Multi-sample batch should have consistent shapes."""
        from lmflow.datasets.multi_modal_dataset import DataCollatorForQwenVL

        collator = DataCollatorForQwenVL(
            processor=mock_processor,
            tokenizer=mock_tokenizer,
            bd_size=None,
            max_length=None,
        )
        instances = [
            {"messages": [
                {"role": "user", "content": [{"type": "text", "text": "Q1"}]},
                {"role": "assistant", "content": [{"type": "text", "text": "A1"}]},
            ]},
            {"messages": [
                {"role": "user", "content": [{"type": "text", "text": "Q2"}]},
                {"role": "assistant", "content": [{"type": "text", "text": "A2"}]},
            ]},
        ]
        batch = collator(instances)

        assert batch["input_ids"].shape[0] == 2
        assert batch["labels"].shape[0] == 2
        assert batch["attention_mask"].shape[0] == 2
        assert batch["input_ids"].shape == batch["labels"].shape


# ============================================================
# 3. Model forward path dispatch tests
# ============================================================

class TestModelForwardDispatch:
    """Test that the VLM model correctly dispatches LLM vs VLM paths."""

    def test_llm_path_when_no_pixel_values(self):
        """#cpu - When pixel_values is None, model should take LLM-only path.
        We test this by checking the code logic, not running full model."""
        # The dispatch logic is at modeling.py line 1829:
        # if pixel_values is None and pixel_values_videos is None: # only train on text
        # This is a structural test - verify the branch exists
        modeling_path = PROJECT_ROOT / "src/lmflow/models/qwen2_5_vl_mdm/modeling.py"
        assert modeling_path.exists()
        content = modeling_path.read_text()
        assert "if pixel_values is None and pixel_values_videos is None:" in content, \
            "LLM-only branch must exist in modeling.py"

    def test_vlm_path_comment_exists(self):
        """#cpu - VLM path should have clear marker."""
        modeling_path = PROJECT_ROOT / "src/lmflow/models/qwen2_5_vl_mdm/modeling.py"
        content = modeling_path.read_text()
        assert "多模态 block diffusion" in content or "multimodal block diffusion" in content.lower(), \
            "VLM path should be clearly marked"

    def test_block_causal_no_dynamic_branch(self):
        """#cpu - block_causal_no_dynamic should have its own branch."""
        modeling_path = PROJECT_ROOT / "src/lmflow/models/qwen2_5_vl_mdm/modeling.py"
        content = modeling_path.read_text()
        assert "self.block_causal_no_dynamic" in content, \
            "block_causal_no_dynamic flag must be used in modeling.py"
        assert "gen_block_causal_mask" in content, \
            "gen_block_causal_mask function must exist for block_causal_no_dynamic path"


# ============================================================
# 4. finetune_qwenvl.py worker cap logic tests
# ============================================================

class TestFinetuneQwenVLWorkerCap:
    """Test the worker count capping logic in finetune_qwenvl.py."""

    def test_worker_cap_applies(self):
        """#cpu - dataloader_num_workers should be capped to max 4."""
        # Simulate the capping logic from finetune_qwenvl.py lines 65-68
        class FakeArgs:
            dataloader_num_workers = 8

        pipeline_args = FakeArgs()
        if hasattr(pipeline_args, "dataloader_num_workers") and pipeline_args.dataloader_num_workers:
            pipeline_args.dataloader_num_workers = max(0, min(4, int(pipeline_args.dataloader_num_workers)))

        assert pipeline_args.dataloader_num_workers == 4, \
            "dataloader_num_workers should be capped to 4"

    def test_worker_cap_preserves_low_values(self):
        """#cpu - Values <= 4 should be preserved."""
        class FakeArgs:
            dataloader_num_workers = 2

        pipeline_args = FakeArgs()
        if hasattr(pipeline_args, "dataloader_num_workers") and pipeline_args.dataloader_num_workers:
            pipeline_args.dataloader_num_workers = max(0, min(4, int(pipeline_args.dataloader_num_workers)))

        assert pipeline_args.dataloader_num_workers == 2

    def test_worker_cap_zero_unchanged(self):
        """#cpu - Zero workers should remain zero (falsy, so cap not applied)."""
        class FakeArgs:
            dataloader_num_workers = 0

        pipeline_args = FakeArgs()
        # The condition `if pipeline_args.dataloader_num_workers:` is False for 0
        if hasattr(pipeline_args, "dataloader_num_workers") and pipeline_args.dataloader_num_workers:
            pipeline_args.dataloader_num_workers = max(0, min(4, int(pipeline_args.dataloader_num_workers)))

        assert pipeline_args.dataloader_num_workers == 0


# ============================================================
# 5. MDM config propagation tests
# ============================================================

class TestMDMConfigPropagation:
    """Test that MDM config is properly set in finetune_qwenvl.py."""

    def test_finetune_qwenvl_sets_mdm_config(self):
        """#cpu - finetune_qwenvl.py should set all MDM config attrs."""
        script_path = PROJECT_ROOT / "examples/finetune_qwenvl.py"
        content = script_path.read_text()

        required_attrs = [
            "bd_size",
            "complementary_mask",
            "always_mask_im_end",
            "flexible_bd_size",
            "use_block_causal_mask",
            "anneal_block_size",
            "block_causal_no_dynamic",
        ]
        for attr in required_attrs:
            assert attr in content, f"finetune_qwenvl.py missing MDM config: {attr}"

    def test_finetune_py_sets_mdm_config(self):
        """#cpu - finetune.py (LLM) should also set MDM config."""
        script_path = PROJECT_ROOT / "examples/finetune.py"
        content = script_path.read_text()

        required_attrs = [
            "bd_size",
            "complementary_mask",
            "always_mask_im_end",
            "use_block_causal_mask",
            "block_causal_no_dynamic",
        ]
        for attr in required_attrs:
            assert attr in content, f"finetune.py missing MDM config: {attr}"

    def test_both_scripts_set_mask_id(self):
        """#cpu - Both scripts should set mask_id for MDM."""
        for script_name in ["finetune.py", "finetune_qwenvl.py"]:
            script_path = PROJECT_ROOT / "examples" / script_name
            content = script_path.read_text()
            assert "mask_id" in content, f"{script_name} must set mask_id"
            assert "|<MASK>|" in content, f"{script_name} must reference |<MASK>| token"


# ============================================================
# 6. CustomMultiModalDataset normalization tests
# ============================================================

class TestCustomMultiModalDatasetNormalization:
    """Test the _normalize_item method of CustomMultiModalDataset."""

    def _normalize_item(self, item):
        """Replicate the normalization logic for testing."""
        image_value = item.get('image')
        if not image_value:
            image_path = item.get('image_path')
            image_paths = item.get('image_paths')
            images = item.get('images')
            if isinstance(image_path, str) and image_path:
                image_value = image_path
            elif isinstance(image_paths, list) and len(image_paths) > 0 and isinstance(image_paths[0], str):
                image_value = image_paths[0]
            elif isinstance(images, list) and len(images) > 0 and isinstance(images[0], str):
                image_value = images[0]

        conv = item.get('conversations')
        if conv is None:
            conv = item.get('messages')

        if not image_value or not isinstance(conv, list) or len(conv) == 0:
            return None

        norm_conv = []
        for m in conv:
            frm = m.get('from')
            if not frm:
                role = m.get('role')
                if role in ('user', 'human'):
                    frm = 'human'
                elif role in ('assistant', 'gpt', 'bot'):
                    frm = 'gpt'
                else:
                    frm = role
            value = m.get('value')
            if value is None:
                value = m.get('text') or m.get('content') or ''
            norm_conv.append({'from': frm, 'value': value})

        if len(norm_conv) == 0:
            return None

        return {'image': image_value, 'conversations': norm_conv}

    def test_standard_format(self):
        """#cpu"""
        item = {
            "image": "test.jpg",
            "conversations": [
                {"from": "human", "value": "What is this?"},
                {"from": "gpt", "value": "An image."},
            ]
        }
        result = self._normalize_item(item)
        assert result is not None
        assert result["image"] == "test.jpg"
        assert len(result["conversations"]) == 2

    def test_image_path_key(self):
        """#cpu"""
        item = {
            "image_path": "alt.jpg",
            "conversations": [{"from": "human", "value": "Q"}]
        }
        result = self._normalize_item(item)
        assert result is not None
        assert result["image"] == "alt.jpg"

    def test_messages_key(self):
        """#cpu - Should accept 'messages' as alternative to 'conversations'."""
        item = {
            "image": "test.jpg",
            "messages": [
                {"role": "user", "text": "Hello"},
                {"role": "assistant", "text": "Hi"},
            ]
        }
        result = self._normalize_item(item)
        assert result is not None
        assert result["conversations"][0]["from"] == "human"
        assert result["conversations"][1]["from"] == "gpt"

    def test_missing_image_returns_none(self):
        """#cpu"""
        item = {
            "conversations": [{"from": "human", "value": "Q"}]
        }
        result = self._normalize_item(item)
        assert result is None

    def test_empty_conversations_returns_none(self):
        """#cpu"""
        item = {"image": "test.jpg", "conversations": []}
        result = self._normalize_item(item)
        assert result is None


# ============================================================
# 7. Profile script existence test
# ============================================================

class TestProfileScript:
    """Test that profiling script exists and is valid."""

    def test_profile_script_exists(self):
        """#cpu"""
        script_path = PROJECT_ROOT / "chengyue_script/dvlm/utils/profile_vlm_vs_llm.py"
        assert script_path.exists(), "profile_vlm_vs_llm.py should exist"

    def test_profile_script_syntax(self):
        """#cpu - Profile script should be valid Python."""
        script_path = PROJECT_ROOT / "chengyue_script/dvlm/utils/profile_vlm_vs_llm.py"
        result = subprocess.run(
            [sys.executable, "-m", "py_compile", str(script_path)],
            capture_output=True, text=True
        )
        assert result.returncode == 0, f"Syntax error in profile script: {result.stderr}"


if __name__ == "__main__":
    pytest.main([__file__, "-v", "--tb=short"])
