"""Instrument discovered existing unit tests and add them to Harbor task directories.

Uses codeflash's inject_profiling_into_existing_test() to create both behavioral
and performance instrumented copies of existing tests.

Must run inside Docker where both inference and codeflash are installed.

Usage:
    python rl_env/instrument_existing_tests.py \
        --existing-tests-map rl_env/existing_tests_map.json \
        --tasks-dir rl_env/codeflash-perf-tasks
"""

from __future__ import annotations

import argparse
import json
import logging
import re
import subprocess
import tempfile
from pathlib import Path

logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
log = logging.getLogger(__name__)

DOCKER_IMAGE = "codeflash-inference-base:latest"


def sanitize_task_name(function_name: str, file_path: str) -> str:
    """Must match task_generator.py's sanitize_task_name."""
    grandparent = Path(file_path).parent.name
    name = f"{grandparent}-{function_name}".replace(".", "-")
    name = re.sub(r"[^a-zA-Z0-9_-]", "-", name)
    name = re.sub(r"-+", "-", name).strip("-").lower()
    return name[:80]


# This script runs INSIDE Docker to instrument one test file
DOCKER_INSTRUMENT_SCRIPT = """
import json
import sys
from pathlib import Path

data = json.loads(sys.argv[1])
function_name = data["function_name"]
file_path = data["file_path"]
test_file = data["test_file"]
call_positions_raw = data["call_positions"]

project_root = Path("/workspace/inference")
test_path = project_root / test_file

if not test_path.exists():
    print(json.dumps({"error": f"Test file not found: {test_path}", "behavior": None, "perf": None}))
    sys.exit(0)

from codeflash.code_utils.instrument_existing_tests import inject_profiling_into_existing_test
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
from codeflash.models.models import CodePosition, FunctionParent, TestingMode

short_name = function_name.rsplit(".", maxsplit=1)[-1]
parents = []
if "." in function_name:
    class_name = function_name.rsplit(".", maxsplit=1)[0]
    parents = [FunctionParent(name=class_name, type="ClassDef")]

fto = FunctionToOptimize(
    function_name=short_name,
    file_path=project_root / file_path,
    parents=parents,
    language="python",
)

# Build CodePosition objects from the discovery map
# If we don't have exact positions, use a broad match (line 0, col 0)
positions = []
for pos in call_positions_raw:
    positions.append(CodePosition(line_no=pos.get("line_no", 0), col_no=pos.get("col_no", 0)))

# If no positions were provided, try to find them via Jedi
if not positions:
    try:
        import jedi
        script = jedi.Script(path=str(test_path), project=jedi.Project(path=str(project_root)))
        names = script.get_names(all_scopes=True, references=True)
        for n in names:
            if n.name == short_name:
                defs = n.goto(follow_imports=True, follow_builtin_imports=False)
                for d in defs:
                    if d.module_path and str(fto.file_path).endswith(str(d.module_path)):
                        positions.append(CodePosition(line_no=n.line, col_no=n.column))
    except Exception:
        pass

if not positions:
    print(json.dumps({"error": "No call positions found", "behavior": None, "perf": None}))
    sys.exit(0)

results = {}
for mode_name, mode in [("behavior", TestingMode.BEHAVIOR), ("perf", TestingMode.PERFORMANCE)]:
    try:
        success, code = inject_profiling_into_existing_test(
            test_path=test_path,
            call_positions=positions,
            function_to_optimize=fto,
            tests_project_root=project_root,
            mode=mode,
        )
        results[mode_name] = code if success else None
    except Exception as e:
        results[mode_name] = None
        results.setdefault("errors", []).append(f"{mode_name}: {e}")

print(json.dumps(results))
"""


def instrument_test_via_docker(
    function_name: str, file_path: str, test_file: str, call_positions: list[dict]
) -> tuple[str | None, str | None]:
    """Instrument a single test file inside Docker. Returns (behavior_code, perf_code)."""
    data = json.dumps(
        {
            "function_name": function_name,
            "file_path": file_path,
            "test_file": test_file,
            "call_positions": call_positions,
        }
    )

    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False, encoding="utf-8") as f:
        f.write(DOCKER_INSTRUMENT_SCRIPT)
        script_path = Path(f.name)

    result = subprocess.run(
        [
            "docker",
            "run",
            "--rm",
            "-v",
            f"{script_path}:/tmp/instrument.py:ro",
            DOCKER_IMAGE,
            "python",
            "/tmp/instrument.py",
            data,
        ],
        capture_output=True,
        text=True,
        encoding="utf-8",
        timeout=120,
    )

    script_path.unlink(missing_ok=True)

    if result.returncode != 0:
        log.debug(f"  Docker instrument failed: {result.stderr[-300:]}")
        return None, None

    for line in reversed(result.stdout.strip().splitlines()):
        if line.strip().startswith("{"):
            try:
                parsed = json.loads(line)
                return parsed.get("behavior"), parsed.get("perf")
            except json.JSONDecodeError:
                continue

    return None, None


def main() -> None:
    parser = argparse.ArgumentParser(description="Instrument existing tests and add to Harbor tasks")
    parser.add_argument("--existing-tests-map", type=Path, default=Path("rl_env/existing_tests_map.json"))
    parser.add_argument("--tasks-dir", type=Path, default=Path("rl_env/codeflash-perf-tasks"))
    args = parser.parse_args()

    test_map = json.loads(args.existing_tests_map.read_text(encoding="utf-8"))
    entries_with_tests = [e for e in test_map if e.get("existing_tests")]

    log.info(f"Found {len(entries_with_tests)} tasks with existing tests to instrument")

    stats = {"instrumented": 0, "failed": 0, "test_files_added": 0}

    for entry in entries_with_tests:
        func_name = entry["function_name"]
        file_path = entry["file_path"]
        task_name = sanitize_task_name(func_name, file_path)
        task_dir = args.tasks_dir / task_name
        test_dir = task_dir / "tests" / "codeflash_eval"

        if not task_dir.exists():
            log.debug(f"  No task dir for {func_name}")
            continue

        log.info(f"  {func_name}: {len(entry['existing_tests'])} existing tests")

        # Group tests by file
        by_file: dict[str, list[dict]] = {}
        for test in entry["existing_tests"]:
            by_file.setdefault(test["test_file"], []).append(test)

        for test_file, tests in by_file.items():
            # Build call positions (we don't have exact positions from discovery, pass empty)
            call_positions: list[dict] = []

            behavior_code, perf_code = instrument_test_via_docker(func_name, file_path, test_file, call_positions)

            if behavior_code:
                short = func_name.rsplit(".", maxsplit=1)[-1]
                test_stem = Path(test_file).stem
                behavior_path = test_dir / f"{test_stem}__{short}__existing_behaviorinstrumented.py"
                behavior_path.write_text(behavior_code, encoding="utf-8")
                stats["test_files_added"] += 1

            if perf_code:
                perf_path = test_dir / f"{test_stem}__{short}__existing_perfinstrumented.py"
                perf_path.write_text(perf_code, encoding="utf-8")
                stats["test_files_added"] += 1

            if behavior_code or perf_code:
                stats["instrumented"] += 1
            else:
                stats["failed"] += 1
                log.warning(f"    Failed to instrument {test_file}")

    log.info("=" * 60)
    log.info("INSTRUMENTATION RESULTS:")
    log.info(f"  Instrumented: {stats['instrumented']} test files")
    log.info(f"  Failed: {stats['failed']} test files")
    log.info(f"  Total files added: {stats['test_files_added']}")


if __name__ == "__main__":
    main()
