# Prompt: 弄几个unittest验证一下 code gen eval integration
"""
Unit tests for the code generation evaluation integration:
  - ChartMimic dataset class mapping in generate_vlmeval_config.py
  - Third-party task routing in run_all_benchmarks.sh
  - Design2Code eval_design2code.py HTML extraction logic
  - HumanEval/MBPP eval_humaneval_mbpp.py argument validation
  - Submodule presence and file structure
"""

import importlib
import json
import os
import subprocess
import sys
from pathlib import Path

import pytest

# ---------------------------------------------------------------------------
# Resolve paths
# ---------------------------------------------------------------------------
SCRIPT_DIR = Path(__file__).resolve().parent
EVAL_DIR = SCRIPT_DIR.parent  # chengyue_script/dvlm/evaluation
VLMEVAL_DIR = EVAL_DIR / "vlmeval"
DESIGN2CODE_DIR = EVAL_DIR / "design2code"
CODE_EVAL_DIR = EVAL_DIR / "code_eval"
PROJECT_ROOT = EVAL_DIR.parent.parent.parent  # LMFlow root

# Add paths for imports
sys.path.insert(0, str(VLMEVAL_DIR))


# ===========================================================================
# 1. generate_vlmeval_config.py – DATASET_CLASS_MAP tests
# ===========================================================================
class TestDatasetClassMap:
    """Verify ChartMimic entries exist and map correctly."""

    @pytest.fixture(autouse=True)
    def _load_module(self):
        spec = importlib.util.spec_from_file_location(
            "generate_vlmeval_config",
            str(VLMEVAL_DIR / "generate_vlmeval_config.py"),
        )
        self.mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(self.mod)

    @pytest.mark.parametrize(
        "dataset_name,expected_class",
        [
            ("ChartMimic_v2_direct_600", "ChartMimic"),
            ("ChartMimic_v2_customized_600", "ChartMimic"),
            ("ChartMimic_v2_direct", "ChartMimic"),
            ("ChartMimic_v2_customized", "ChartMimic"),
            ("ChartMimic_v2_direct_1800", "ChartMimic"),
            ("ChartMimic_v2_customized_1800", "ChartMimic"),
            ("ChartMimic_v1_direct", "ChartMimic"),
            ("ChartMimic_v1_customized", "ChartMimic"),
        ],
    )
    def test_chartmimic_in_map(self, dataset_name, expected_class):
        assert dataset_name in self.mod.DATASET_CLASS_MAP
        assert self.mod.DATASET_CLASS_MAP[dataset_name] == expected_class

    @pytest.mark.parametrize(
        "dataset_name,expected_class",
        [
            ("ChartMimic_v2_direct_600", "ChartMimic"),
            ("ChartMimic_v2_customized_600", "ChartMimic"),
            # Existing datasets should still work
            ("GQA_TestDev_Balanced", "ImageVQADataset"),
            ("MMMU_DEV_VAL", "MMMUDataset"),
            ("MathVista_MINI", "MathVista"),
            ("POPE", "ImageYORNDataset"),
        ],
    )
    def test_get_dataset_class(self, dataset_name, expected_class):
        assert self.mod.get_dataset_class(dataset_name) == expected_class

    def test_unknown_dataset_returns_default(self):
        assert self.mod.get_dataset_class("SomeUnknownBench") == "ImageMCQDataset"

    def test_config_generation_chartmimic(self, tmp_path):
        """End-to-end: generate config JSON with ChartMimic datasets (via subprocess)."""
        cfg_path = tmp_path / "config.json"
        env = {
            "MODEL_PATH_ABS": "/tmp/dummy",
            "PROCESSOR_PATH": "Qwen/Qwen2.5-VL-3B-Instruct",
            "DATASETS": "ChartMimic_v2_direct_600 ChartMimic_v2_customized_600",
            "MODEL_KEY": "test",
            "CFG_PATH": str(cfg_path),
            "AUTO_DETECT_CLASS": "1",
            "PATH": os.environ.get("PATH", ""),
            "HOME": os.environ.get("HOME", ""),
        }
        result = subprocess.run(
            [sys.executable, str(VLMEVAL_DIR / "generate_vlmeval_config.py")],
            env=env,
            capture_output=True,
            text=True,
            timeout=30,
        )
        assert result.returncode == 0, f"stderr: {result.stderr}"
        cfg = json.loads(cfg_path.read_text())
        assert "ChartMimic_v2_direct_600" in cfg["data"]
        assert cfg["data"]["ChartMimic_v2_direct_600"]["class"] == "ChartMimic"
        assert "ChartMimic_v2_customized_600" in cfg["data"]
        assert cfg["data"]["ChartMimic_v2_customized_600"]["class"] == "ChartMimic"


# ===========================================================================
# 2. run_all_benchmarks.sh – third-party routing tests
# ===========================================================================
class TestThirdPartyRouting:
    """Verify the shell script correctly identifies third-party tasks."""

    SHELL_SCRIPT = str(VLMEVAL_DIR / "run_all_benchmarks.sh")

    def _run_is_third_party(self, task: str) -> bool:
        """Inline the is_third_party_task logic and test it via bash."""
        cmd = f"""
        is_third_party_task() {{
            local task=$1
            case "${{task}}" in
                Design2Code|HumanEval|MBPP) return 0 ;;
                *) return 1 ;;
            esac
        }}
        is_third_party_task "{task}" && echo "YES" || echo "NO"
        """
        result = subprocess.run(
            ["bash", "-c", cmd],
            capture_output=True,
            text=True,
            timeout=5,
        )
        return result.stdout.strip() == "YES"

    @pytest.mark.parametrize("task", ["Design2Code", "HumanEval", "MBPP"])
    def test_third_party_tasks_detected(self, task):
        assert self._run_is_third_party(task) is True

    @pytest.mark.parametrize(
        "task",
        [
            "MMMU_DEV_VAL",
            "ChartMimic_v2_direct_600",
            "GQA_TestDev_Balanced",
            "POPE",
        ],
    )
    def test_vlmeval_tasks_not_detected_as_third_party(self, task):
        assert self._run_is_third_party(task) is False

    def test_help_mentions_third_party(self):
        """Help text should mention Design2Code, HumanEval, MBPP."""
        content = Path(self.SHELL_SCRIPT).read_text()
        assert "Design2Code" in content
        assert "HumanEval" in content
        assert "MBPP" in content

    def test_help_mentions_chartmimic(self):
        content = Path(self.SHELL_SCRIPT).read_text()
        assert "ChartMimic_v2_direct_600" in content
        assert "ChartMimic_v2_customized_600" in content

    def test_default_tasks_include_chartmimic(self):
        """DEFAULT_TASKS variable should include ChartMimic."""
        content = Path(self.SHELL_SCRIPT).read_text()
        for line in content.splitlines():
            if line.startswith("DEFAULT_TASKS="):
                assert "ChartMimic_v2_direct_600" in line
                assert "ChartMimic_v2_customized_600" in line
                return
        pytest.fail("DEFAULT_TASKS= line not found")

    def test_chartmimic_judge_is_gpt4o(self):
        """ChartMimic tasks should route to gpt-4o judge."""
        content = Path(self.SHELL_SCRIPT).read_text()
        assert "ChartMimic*)" in content
        lines = content.splitlines()
        for i, line in enumerate(lines):
            if "ChartMimic*)" in line:
                for j in range(i + 1, min(i + 3, len(lines))):
                    if "gpt-4o" in lines[j]:
                        return
        pytest.fail("ChartMimic judge routing to gpt-4o not found")

    def test_run_third_party_task_function_exists(self):
        """run_third_party_task function should be defined in the script."""
        content = Path(self.SHELL_SCRIPT).read_text()
        assert "run_third_party_task()" in content

    def test_is_third_party_task_function_exists(self):
        """is_third_party_task function should be defined in the script."""
        content = Path(self.SHELL_SCRIPT).read_text()
        assert "is_third_party_task()" in content


# ===========================================================================
# 3. Design2Code – HTML extraction logic
# ===========================================================================
class TestDesign2CodeHTMLExtraction:
    """Test HTML extraction from model output (replicated logic, no heavy imports)."""

    @staticmethod
    def extract_html(raw: str) -> str:
        """Replicate the extraction logic from eval_design2code.py."""
        text = raw
        if "```html" in text:
            text = text.split("```html", 1)[1]
            if "```" in text:
                text = text.split("```", 1)[0]
        elif "```" in text:
            text = text.split("```", 1)[1]
            if "```" in text:
                text = text.split("```", 1)[0]
        return text.strip()

    def test_html_code_block(self):
        raw = "```html\n<html><body>Hello</body></html>\n```"
        assert self.extract_html(raw) == "<html><body>Hello</body></html>"

    def test_plain_html_passthrough(self):
        text = "<html><body>Test</body></html>"
        assert self.extract_html(text) == "<html><body>Test</body></html>"

    def test_generic_code_block(self):
        raw = "```\n<div>content</div>\n```"
        assert self.extract_html(raw) == "<div>content</div>"

    def test_html_with_preamble(self):
        raw = "Here is the HTML:\n```html\n<!DOCTYPE html>\n<html></html>\n```\nDone."
        result = self.extract_html(raw)
        assert result.startswith("<!DOCTYPE html>")
        assert result.endswith("<html></html>")

    def test_multiline_html(self):
        raw = "```html\n<html>\n  <body>\n    <p>Para</p>\n  </body>\n</html>\n```"
        result = self.extract_html(raw)
        assert "<p>Para</p>" in result
        assert result.startswith("<html>")


# ===========================================================================
# 4. HumanEval/MBPP – argparse validation (via subprocess, lightweight)
# ===========================================================================
class TestHumanEvalMBPPArgs:
    """Test argument parsing without loading model/GPU code."""

    SCRIPT = str(CODE_EVAL_DIR / "eval_humaneval_mbpp.py")

    def test_task_required(self):
        """--task is required, should fail without it."""
        result = subprocess.run(
            [sys.executable, self.SCRIPT, "--model_path", "/tmp/dummy"],
            capture_output=True,
            text=True,
            timeout=30,
        )
        assert result.returncode != 0
        assert "required" in result.stderr.lower() or "error" in result.stderr.lower()

    def test_invalid_task_rejected(self):
        """Invalid task name should be rejected by choices=."""
        result = subprocess.run(
            [
                sys.executable,
                self.SCRIPT,
                "--model_path",
                "/tmp/dummy",
                "--task",
                "invalid_bench",
            ],
            capture_output=True,
            text=True,
            timeout=30,
        )
        assert result.returncode != 0
        assert "invalid choice" in result.stderr.lower()


# ===========================================================================
# 5. Submodule presence tests
# ===========================================================================
class TestSubmodules:
    """Verify third-party submodules are registered."""

    def test_gitmodules_has_design2code(self):
        gitmodules = (PROJECT_ROOT / ".gitmodules").read_text()
        assert "third_party/Design2Code" in gitmodules

    def test_gitmodules_has_bigcode_eval(self):
        gitmodules = (PROJECT_ROOT / ".gitmodules").read_text()
        assert "third_party/bigcode-eval" in gitmodules

    def test_gitmodules_has_vlmevalkit(self):
        gitmodules = (PROJECT_ROOT / ".gitmodules").read_text()
        assert "third_party/VLMEvalKit" in gitmodules

    def test_chartmimic_file_exists_in_vlmevalkit(self):
        chartmimic_py = (
            PROJECT_ROOT
            / "third_party"
            / "VLMEvalKit"
            / "vlmeval"
            / "dataset"
            / "chartmimic.py"
        )
        assert chartmimic_py.exists(), f"ChartMimic module not found at {chartmimic_py}"

    def test_design2code_metrics_exist(self):
        metrics_py = (
            PROJECT_ROOT
            / "third_party"
            / "Design2Code"
            / "Design2Code"
            / "metrics"
            / "multi_processing_eval.py"
        )
        assert metrics_py.exists(), f"Design2Code metrics not found at {metrics_py}"

    def test_bigcode_eval_main_exists(self):
        main_py = PROJECT_ROOT / "third_party" / "bigcode-eval" / "main.py"
        assert main_py.exists(), f"bigcode-eval main.py not found at {main_py}"


# ===========================================================================
# 6. Eval script file structure & syntax tests
# ===========================================================================
class TestEvalFileStructure:
    """Verify all evaluation scripts exist, are executable, and have valid syntax."""

    def test_design2code_py_exists(self):
        assert (DESIGN2CODE_DIR / "eval_design2code.py").exists()

    def test_design2code_sh_exists_and_executable(self):
        sh = DESIGN2CODE_DIR / "eval_design2code.sh"
        assert sh.exists()
        assert os.access(sh, os.X_OK), "eval_design2code.sh is not executable"

    def test_humaneval_mbpp_py_exists(self):
        assert (CODE_EVAL_DIR / "eval_humaneval_mbpp.py").exists()

    def test_humaneval_mbpp_sh_exists_and_executable(self):
        sh = CODE_EVAL_DIR / "eval_humaneval_mbpp.sh"
        assert sh.exists()
        assert os.access(sh, os.X_OK), "eval_humaneval_mbpp.sh is not executable"

    def test_design2code_py_syntax(self):
        result = subprocess.run(
            [sys.executable, "-m", "py_compile", str(DESIGN2CODE_DIR / "eval_design2code.py")],
            capture_output=True,
            text=True,
            timeout=30,
        )
        assert result.returncode == 0, f"Syntax error: {result.stderr}"

    def test_humaneval_mbpp_py_syntax(self):
        result = subprocess.run(
            [sys.executable, "-m", "py_compile", str(CODE_EVAL_DIR / "eval_humaneval_mbpp.py")],
            capture_output=True,
            text=True,
            timeout=30,
        )
        assert result.returncode == 0, f"Syntax error: {result.stderr}"

    def test_design2code_sh_syntax(self):
        result = subprocess.run(
            ["bash", "-n", str(DESIGN2CODE_DIR / "eval_design2code.sh")],
            capture_output=True,
            text=True,
            timeout=10,
        )
        assert result.returncode == 0, f"Syntax error: {result.stderr}"

    def test_humaneval_mbpp_sh_syntax(self):
        result = subprocess.run(
            ["bash", "-n", str(CODE_EVAL_DIR / "eval_humaneval_mbpp.sh")],
            capture_output=True,
            text=True,
            timeout=10,
        )
        assert result.returncode == 0, f"Syntax error: {result.stderr}"

    def test_run_all_benchmarks_sh_syntax(self):
        result = subprocess.run(
            ["bash", "-n", str(VLMEVAL_DIR / "run_all_benchmarks.sh")],
            capture_output=True,
            text=True,
            timeout=10,
        )
        assert result.returncode == 0, f"Syntax error: {result.stderr}"

    def test_generate_vlmeval_config_py_syntax(self):
        result = subprocess.run(
            [sys.executable, "-m", "py_compile", str(VLMEVAL_DIR / "generate_vlmeval_config.py")],
            capture_output=True,
            text=True,
            timeout=30,
        )
        assert result.returncode == 0, f"Syntax error: {result.stderr}"
