"""DebugBench dataset importer for llmdebug eval framework.

Source: HuggingFace ``Rtian/DebugBench`` (Apache-2.0 license).
~1,414 Python cases with GPT-4-implanted bugs in LeetCode-style problems.

Each case has a difficulty level and bug type taxonomy matching our needs.
"""

from __future__ import annotations

import re
import textwrap
from pathlib import Path
from typing import Any

from evals.category_mapping import infer_snapshot_dependency, map_debugbench_type
from evals.importers.base import (
    count_patch_lines,
    generate_oracle_patch,
    validate_generated_case,
    write_case_files,
)

# ---------------------------------------------------------------------------
# Filter criteria
# ---------------------------------------------------------------------------

MAX_FIX_LINES = 15
MAX_CODE_LINES = 200
DISALLOWED_IMPORTS = re.compile(
    r"^\s*(?:import|from)\s+(?!__future__|collections|itertools|functools|math|re|string|typing|dataclasses|enum|operator|copy|heapq|bisect|random|sys|os|io|json|abc|contextlib|decimal|fractions|statistics|unittest)"
    r"(?!\.)",
    re.MULTILINE,
)


def _has_external_deps(code: str) -> bool:
    """Check if code imports anything beyond stdlib (conservative)."""
    return bool(DISALLOWED_IMPORTS.search(code))


def _is_deterministic(code: str) -> bool:
    """Heuristic: no network, no real randomness, no timing."""
    risky = ["requests.", "urllib.", "http.", "socket.", "time.sleep", "random.random"]
    lower = code.lower()
    return not any(r in lower for r in risky)


# ---------------------------------------------------------------------------
# Difficulty mapping
# ---------------------------------------------------------------------------

DIFFICULTY_MAP = {
    "easy": "easy",
    "medium": "medium",
    "hard": "hard",
}


# ---------------------------------------------------------------------------
# Test generation
# ---------------------------------------------------------------------------


def _generate_test(buggy_code: str, test_list: list[str] | str | None) -> str:
    """Generate test_buggy.py from DebugBench test cases."""
    tests: list[str] = []
    if isinstance(test_list, list):
        tests = test_list
    elif isinstance(test_list, str):
        tests = [test_list]

    if not tests:
        return textwrap.dedent("""\
            from buggy import *

            def test_basic():
                # Placeholder: DebugBench case needs manual test
                pass
        """)

    lines = ["from buggy import *", "", ""]
    for i, test_code in enumerate(tests):
        test_code = test_code.strip()
        if not test_code:
            continue
        fn_name = f"test_case_{i}"
        # If the test is already a function definition, include as-is.
        if test_code.startswith("def "):
            lines.append(test_code)
        else:
            # Wrap assertion-style tests.
            lines.append(f"def {fn_name}():")
            for line in test_code.splitlines():
                lines.append(f"    {line}")
        lines.append("")
        lines.append("")

    return "\n".join(lines).rstrip() + "\n"


# ---------------------------------------------------------------------------
# Infer exception type from buggy code
# ---------------------------------------------------------------------------


def _infer_exception(bug_type: str, buggy_code: str) -> str:
    """Best-effort exception type inference from bug type and code."""
    bug_lower = bug_type.lower()
    if "syntax" in bug_lower:
        return "SyntaxError"
    if "reference" in bug_lower:
        return "NameError"
    if "type" in bug_lower:
        return "TypeError"
    # Default for logic errors — tests typically raise AssertionError.
    return "AssertionError"


# ---------------------------------------------------------------------------
# Main import function
# ---------------------------------------------------------------------------


def import_debugbench(
    output_dir: Path,
    *,
    max_cases: int = 40,
    validate: bool = True,
    dry_run: bool = False,
) -> list[dict[str, Any]]:
    """Import cases from DebugBench dataset.

    Returns list of imported case metadata dicts (or preview dicts if dry_run).
    """
    try:
        from datasets import load_dataset  # type: ignore[import-untyped]
    except ImportError:
        raise ImportError(
            "The 'datasets' package is required for DebugBench import. "
            "Install with: pip install datasets"
        ) from None

    ds = load_dataset("Rtian/DebugBench", split="train")

    imported: list[dict[str, Any]] = []
    skipped = 0

    for idx, row in enumerate(ds):
        if len(imported) >= max_cases:
            break

        # Extract fields.
        language = str(row.get("language", "")).lower()
        if language != "python":
            continue

        buggy_code = str(row.get("buggy_code", ""))
        fixed_code = str(row.get("code", ""))
        bug_type = str(row.get("bug_type", ""))
        difficulty = str(row.get("difficulty", "medium")).lower()
        test_list = row.get("test_list")

        if not buggy_code.strip() or not fixed_code.strip():
            skipped += 1
            continue

        # Apply filters.
        oracle_patch = generate_oracle_patch(buggy_code, fixed_code)
        if not oracle_patch.strip():
            skipped += 1
            continue

        fix_lines = count_patch_lines(oracle_patch)
        if fix_lines > MAX_FIX_LINES:
            skipped += 1
            continue

        if len(buggy_code.splitlines()) > MAX_CODE_LINES:
            skipped += 1
            continue

        if _has_external_deps(buggy_code) or _has_external_deps(fixed_code):
            skipped += 1
            continue

        if not _is_deterministic(buggy_code):
            skipped += 1
            continue

        # Map metadata.
        category = map_debugbench_type(bug_type)
        mapped_difficulty = DIFFICULTY_MAP.get(difficulty, "medium")
        snapshot_dep = infer_snapshot_dependency(category)
        expected_exception = _infer_exception(bug_type, buggy_code)
        case_id = f"db_{idx:04d}_{mapped_difficulty}"

        test_code = _generate_test(buggy_code, test_list)

        metadata = {
            "description": f"DebugBench #{idx}: {bug_type} bug ({difficulty})",
            "category": category,
            "difficulty": mapped_difficulty,
            "snapshot_dependency": snapshot_dep,
            "expected_exception": expected_exception,
            "tags": ["debugbench", bug_type.replace(" ", "_")],
            "python_args": ["-m", "pytest", "-q"],
            "timeout_sec": 30,
            "bug_source": "debugbench",
            "contamination_risk": "medium",
            "expected_fix_lines": fix_lines,
            "source_id": f"debugbench_{idx}",
        }

        case_info = {"case_id": case_id, "metadata": metadata, "fix_lines": fix_lines}

        if dry_run:
            imported.append(case_info)
            continue

        case_dir = output_dir / case_id
        write_case_files(
            case_dir,
            buggy_code=buggy_code,
            test_code=test_code,
            oracle_patch=oracle_patch,
            metadata=metadata,
        )

        if validate:
            ok, msg = validate_generated_case(case_dir, schema_only=True)
            if not ok:
                skipped += 1
                # Clean up invalid case.
                import shutil

                shutil.rmtree(case_dir, ignore_errors=True)
                continue

        imported.append(case_info)

    return imported
