#!/usr/bin/env python3
"""Codeflash RL Environment Evaluator.

Evaluates an agent's code optimization for correctness and performance,
reusing codeflash's production verification and benchmarking pipeline.

Writes reward.json to the output directory in Harbor format.

Usage (inside Docker):
    python /tests/codeflash_eval/evaluate.py \
        --function "_evict_oldest_video" \
        --file-path "inference/core/workflows/core_steps/analytics/detection_event_log/v1.py" \
        --repo-path /workspace/inference \
        --test-dir /tests/codeflash_eval \
        --output-dir /logs/verifier \
        --debug
"""

from __future__ import annotations

import argparse
import ast
import json
import logging
import os
import subprocess
import sys
import tempfile
from pathlib import Path

log = logging.getLogger(__name__)

# Will be set in main() based on --debug flag
DEBUG = False

# --- Codeflash imports (required — installed in Docker base image) ---
from codeflash.models.models import (
    TestFile,
    TestFiles,
    TestResults,
    TestType,
)
from codeflash.verification.comparator import comparator as codeflash_comparator
from codeflash.verification.equivalence import compare_test_results
from codeflash.verification.parse_test_output import parse_test_results
from codeflash.verification.verification_utils import TestConfig



BLOCKED_PYTEST_PLUGINS_BEHAVIOR = ["benchmark", "codspeed", "xdist", "sugar"]
BLOCKED_PYTEST_PLUGINS_PERF = ["codspeed", "cov", "benchmark", "profiling", "xdist", "sugar"]


def comparator(a: object, b: object, superset_obj: bool = False) -> bool:
    return codeflash_comparator(a, b, superset_obj=superset_obj)


# ---------------------------------------------------------------------------
# AST-based function replacement
# ---------------------------------------------------------------------------


def find_function_in_source(source: str, func_name: str) -> tuple[int, int, int] | None:
    try:
        tree = ast.parse(source)
    except SyntaxError:
        return None
    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef):
            if node.name == func_name:
                start = node.lineno - 1
                end = node.end_lineno or start + 1
                if node.decorator_list:
                    start = node.decorator_list[0].lineno - 1
                return (start, end, node.col_offset)
    return None



def git_checkout(repo_path: Path, ref: str) -> bool:
    """Checkout a git ref, forcing clean state. Returns True on success."""
    # Use -f to discard any local changes (e.g., .pyc files, coverage artifacts)
    result = subprocess.run(
        ["git", "checkout", "-f", ref], cwd=repo_path, capture_output=True, text=True, encoding="utf-8",
    )
    if DEBUG:
        log.debug(f"  [git] checkout -f {ref[:12]}: {'OK' if result.returncode == 0 else 'FAILED'}")
        if result.returncode != 0:
            log.debug(f"  [git] stderr: {result.stderr[:200]}")
    return result.returncode == 0


def git_get_head(repo_path: Path) -> str:
    """Get current HEAD commit SHA."""
    result = subprocess.run(
        ["git", "rev-parse", "HEAD"], cwd=repo_path, capture_output=True, text=True, encoding="utf-8",
    )
    return result.stdout.strip()


def git_get_parent(repo_path: Path, ref: str) -> str:
    """Get parent commit SHA of a ref."""
    result = subprocess.run(
        ["git", "rev-parse", f"{ref}~1"], cwd=repo_path, capture_output=True, text=True, encoding="utf-8",
    )
    return result.stdout.strip()


# ---------------------------------------------------------------------------
# Test execution
# ---------------------------------------------------------------------------


def build_pytest_cmd(
    test_files: list[Path], xml_result_path: Path, *, mode: str = "behavior", timeout: int | None = None
) -> list[str]:
    cmd = [sys.executable, "-m", "pytest"]
    cmd.extend(["--capture=tee-sys", "-q"])
    cmd.append("--codeflash_loops_scope=session")
    cmd.extend([f"--junitxml={xml_result_path}", "-o", "junit_logging=all"])

    if mode == "behavior":
        cmd.extend(["--codeflash_min_loops=1", "--codeflash_max_loops=1", "--codeflash_seconds=30"])
        blocked = BLOCKED_PYTEST_PLUGINS_BEHAVIOR
    else:
        cmd.extend(
            [
                "--codeflash_min_loops=5",
                "--codeflash_max_loops=100000",
                "--codeflash_seconds=10",
                "--codeflash_stability_check=true",
            ]
        )
        blocked = BLOCKED_PYTEST_PLUGINS_PERF

    if timeout is not None:
        cmd.append(f"--timeout={timeout}")
    for plugin in blocked:
        cmd.extend(["-p", f"no:{plugin}"])
    cmd.extend([str(f) for f in test_files])
    return cmd


def run_tests(
    test_files: list[Path],
    repo_path: Path,
    tmp_dir: Path,
    iteration: int,
    *,
    mode: str = "behavior",
    timeout: int = 300,
) -> tuple[subprocess.CompletedProcess[str], Path, Path]:
    xml_path = tmp_dir / f"pytest_results_{iteration}.xml"
    bin_path = tmp_dir / f"test_return_values_{iteration}.bin"
    cmd = build_pytest_cmd(test_files, xml_path, mode=mode)

    env = os.environ.copy()
    env["PYTHONPATH"] = str(repo_path)
    env["PYTEST_PLUGINS"] = "codeflash.verification.pytest_plugin"
    env["CODEFLASH_TEST_ITERATION"] = str(iteration)
    env["CODEFLASH_LOOP_INDEX"] = "1"
    env["CODEFLASH_TRACER_DISABLE"] = "1"

    if DEBUG:
        log.debug(f"  [pytest] cmd: {' '.join(cmd)}")
        log.debug(f"  [pytest] cwd: {repo_path}")
        log.debug(f"  [pytest] PYTHONPATH: {env['PYTHONPATH']}")
        log.debug(f"  [pytest] CODEFLASH_TEST_ITERATION: {iteration}")
        log.debug(f"  [pytest] test files: {[f.name for f in test_files]}")
        log.debug(f"  [pytest] bin_path will be: {bin_path}")

    log.info(f"  Running {mode} tests (iteration={iteration}, {len(test_files)} files)...")
    result = subprocess.run(
        cmd, cwd=repo_path, capture_output=True, text=True, encoding="utf-8", timeout=timeout, env=env
    )

    if DEBUG:
        log.debug(f"  [pytest] returncode: {result.returncode}")
        log.debug(f"  [pytest] stdout ({len(result.stdout)} chars):\n{result.stdout[-2000:]}")
        if result.stderr:
            log.debug(f"  [pytest] stderr ({len(result.stderr)} chars):\n{result.stderr[-1000:]}")
        log.debug(f"  [pytest] bin file exists: {bin_path.exists()}")
        if bin_path.exists():
            log.debug(f"  [pytest] bin file size: {bin_path.stat().st_size} bytes")
        log.debug(f"  [pytest] xml file exists: {xml_path.exists()}")

    if result.returncode not in (0, 1):
        log.warning(f"  pytest returned code {result.returncode}")
        if result.stderr:
            log.warning(f"  stderr (last 500): {result.stderr[-500:]}")
    return result, xml_path, bin_path


# ---------------------------------------------------------------------------
# Result parsing
# ---------------------------------------------------------------------------



# ---------------------------------------------------------------------------
# Runtime calculation
# ---------------------------------------------------------------------------


def compute_per_test_speedup(original: TestResults, candidate: TestResults) -> dict[str, float]:
    """Per-test-function speedup, aggregated across all invocations within each test.

    Groups invocations by (test_module, test_class, test_function) — the call site level.
    For each test function, sums the min-across-loops runtime of all its invocations.
    This matches how codeflash reports per-test timings in PR comments.
    """
    orig_data = original.usable_runtime_data_by_test_case()
    cand_data = candidate.usable_runtime_data_by_test_case()

    # Aggregate: sum of min(loop runtimes) per invocation, grouped by test function
    orig_by_test: dict[str, int] = {}
    for inv_id, runtimes in orig_data.items():
        test_file = inv_id.test_module_path.rsplit(".", maxsplit=1)[-1] if "." in inv_id.test_module_path else inv_id.test_module_path
        cls = f"{inv_id.test_class_name}." if inv_id.test_class_name else ""
        key = f"{test_file}::{cls}{inv_id.test_function_name}"
        orig_by_test[key] = orig_by_test.get(key, 0) + min(runtimes)

    cand_by_test: dict[str, int] = {}
    for inv_id, runtimes in cand_data.items():
        test_file = inv_id.test_module_path.rsplit(".", maxsplit=1)[-1] if "." in inv_id.test_module_path else inv_id.test_module_path
        cls = f"{inv_id.test_class_name}." if inv_id.test_class_name else ""
        key = f"{test_file}::{cls}{inv_id.test_function_name}"
        cand_by_test[key] = cand_by_test.get(key, 0) + min(runtimes)

    speedups: dict[str, float] = {}
    for key in set(orig_by_test) & set(cand_by_test):
        if cand_by_test[key] > 0:
            speedups[key] = round(orig_by_test[key] / cand_by_test[key], 4)

    return speedups


def _extract_function(source: str, func_name: str) -> str | None:
    """Extract a function's source code from a file."""
    info = find_function_in_source(source, func_name)
    if not info:
        return None
    start, end, _ = info
    lines = source.splitlines(keepends=True)
    return "".join(lines[start:end])


# ---------------------------------------------------------------------------
# Quality rating
# ---------------------------------------------------------------------------


def get_optimization_quality(
    original_code: str,
    optimized_code: str,
    file_path: str,
    speedup: float,
    original_runtime_ns: int = 0,
    candidate_runtime_ns: int = 0,
) -> str:
    """Call aiservice /optimization_review endpoint directly. Returns 'high'/'medium'/'low'/''."""
    try:
        from codeflash.api.aiservice import AiServiceClient
        from codeflash.code_utils.time_utils import humanize_runtime
        from codeflash.version import __version__ as codeflash_version

        client = AiServiceClient()
        # speedup ratio is original/candidate (e.g. 2.0 = 2x faster)
        # The API expects percentage improvement: (speedup - 1) * 100
        speedup_pct = (speedup - 1) * 100
        payload = {
            "original_code": original_code,
            "optimized_code": optimized_code,
            "explanation": f"The optimized code is {speedup:.2f}x faster ({speedup_pct:.0f}% improvement).",
            "existing_tests": "",
            "generated_tests": "",
            "trace_id": "",
            "coverage_message": "",
            "replay_tests": "",
            "speedup": f"{speedup_pct:.2f}%",
            "loop_count": 100,
            "benchmark_details": None,
            "optimized_runtime": humanize_runtime(candidate_runtime_ns) if candidate_runtime_ns else "",
            "original_runtime": humanize_runtime(original_runtime_ns) if original_runtime_ns else "",
            "codeflash_version": codeflash_version,
            "calling_fn_details": "",
            "language": "python",
            "language_version": "",
            "call_sequence": client.get_next_sequence(),
        }
        response = client.make_ai_service_request("/optimization_review", payload=payload, timeout=30)
        if response.status_code == 200:
            data = response.json()
            return data.get("review", "")
        if DEBUG:
            log.debug(f"  [quality] API returned {response.status_code}: {response.text[:200]}")
        return ""
    except Exception as e:
        if DEBUG:
            log.debug(f"  [quality] aiservice call failed: {e}")
        return ""


# ---------------------------------------------------------------------------
# Test file discovery and preparation
# ---------------------------------------------------------------------------


def discover_test_files(test_dir: Path) -> tuple[list[Path], list[Path]]:
    behavioral = sorted(test_dir.glob("*behaviorinstrumented*.py"))
    performance = sorted(test_dir.glob("*perfinstrumented*.py"))
    if DEBUG:
        log.debug(f"  [discover] behavioral: {[f.name for f in behavioral]}")
        log.debug(f"  [discover] performance: {[f.name for f in performance]}")
    return behavioral, performance



def build_test_files_and_config(
    test_paths: list[Path],
    perf_paths: list[Path],
    repo_path: Path,
) -> tuple[TestFiles, TestConfig]:
    """Build TestFiles and TestConfig for codeflash's parse_test_results."""
    test_file_list = []
    for tf in test_paths:
        # Find matching perf test by stem (behavioral_0 <-> perf_0)
        matching_perf = None
        for pf in perf_paths:
            if pf.stem.replace("perfinstrumented", "") == tf.stem.replace("behaviorinstrumented", ""):
                matching_perf = pf
                break

        test_file_list.append(TestFile(
            instrumented_behavior_file_path=tf,
            benchmarking_file_path=matching_perf or tf,
            test_type=TestType.GENERATED_REGRESSION,
        ))

    test_files = TestFiles(test_files=test_file_list)
    test_config = TestConfig(
        tests_root=repo_path / "tests",
        project_root_path=repo_path,
        tests_project_rootdir=repo_path,
    )

    if DEBUG:
        log.debug(f"  [test_files] {len(test_file_list)} TestFile entries:")
        for tf in test_file_list:
            log.debug(f"    behavior={tf.instrumented_behavior_file_path}")

    return test_files, test_config


def run_parse_test_results(
    xml_path: Path,
    test_files: TestFiles,
    test_config: TestConfig,
    iteration: int,
    function_name: str,
    source_file: Path,
    run_result: subprocess.CompletedProcess | None = None,
) -> TestResults:
    """Call codeflash's parse_test_results with proper arguments."""
    from codeflash.languages.current import set_current_language

    set_current_language("python")

    results, _coverage = parse_test_results(
        test_xml_path=xml_path,
        test_files=test_files,
        test_config=test_config,
        optimization_iteration=iteration,
        function_name=function_name,
        source_file=source_file,
        coverage_database_file=None,
        coverage_config_file=None,
        run_result=run_result,
    )

    if DEBUG:
        log.debug(f"  [parse] iteration={iteration}: {len(results)} invocations, "
                  f"{results.number_of_loops()} loops")

    return results


def get_codeflash_tmp_dir() -> Path:
    """Get codeflash's tmp dir (same one that get_run_tmp_file uses).

    This ensures generated tests (.bin) and existing tests (.sqlite) write to the same place
    so parse_test_results() can find both.
    """
    try:
        from codeflash.code_utils.codeflash_wrap_decorator import get_run_tmp_file
    except ImportError:
        from codeflash.code_utils.code_utils import get_run_tmp_file

    return get_run_tmp_file(Path()).parent


def place_test_files(test_dir: Path, repo_path: Path, pattern: str) -> list[Path]:
    """Place test files into the repo's test directory, preserving filenames.

    Replaces the {codeflash_run_tmp_dir_client_side} placeholder with codeflash's
    actual tmp dir so bin/sqlite output goes where parse_test_results expects.
    """
    cf_tmp = get_codeflash_tmp_dir()
    dest_dir = repo_path / "tests" / "codeflash_generated"
    dest_dir.mkdir(parents=True, exist_ok=True)
    placed: list[Path] = []

    for test_file in sorted(test_dir.glob(pattern)):
        content = test_file.read_text(encoding="utf-8")
        content = content.replace("{codeflash_run_tmp_dir_client_side}", str(cf_tmp))

        target = dest_dir / test_file.name
        target.write_text(content, encoding="utf-8")
        placed.append(target)

        if DEBUG:
            log.debug(f"  [place] {test_file.name} → {target.relative_to(repo_path)}")

    if DEBUG:
        log.debug(f"  [place] placed {len(placed)} files matching {pattern}")

    return placed


# ---------------------------------------------------------------------------
# Reward output
# ---------------------------------------------------------------------------


def write_reward(
    output_dir: Path,
    *,
    correct: float,
    speedup: float,
    speedup_per_test: dict[str, float],
    tests_passed: int,
    tests_total: int,
    optimization_quality: str = "",
    coverage_pct: float | None = None,
) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)

    # Harbor reward.json: single key-value pair for the metric Harbor tracks
    reward_for_harbor = {"speedup": speedup if correct == 1.0 else 0.0}
    (output_dir / "reward.json").write_text(json.dumps(reward_for_harbor, indent=2), encoding="utf-8")

    # Full metrics in a separate file for detailed analysis
    details = {
        "correct": correct,
        "speedup": speedup if correct == 1.0 else 0.0,
        "speedup_per_test_case": speedup_per_test if correct == 1.0 else {},
        "tests_passed": tests_passed,
        "tests_total": tests_total,
        "coverage_pct": coverage_pct,
        "optimization_quality": optimization_quality,
    }
    (output_dir / "details.json").write_text(json.dumps(details, indent=2), encoding="utf-8")

    log.info(f"  Wrote reward.json: correct={correct} speedup={speedup} passed={tests_passed}/{tests_total}")


# ---------------------------------------------------------------------------
# Coverage measurement
# ---------------------------------------------------------------------------


def compute_coverage(
    behavioral_tests: list[Path],
    repo_path: Path,
    file_path: str,
    function_name: str,
    tmp_dir: Path,
    original_commit: str,
    candidate_commit: str,
) -> float | None:
    """Compute function-level coverage using codeflash's coverage pipeline.

    Uses coverage.py with branch tracking (SQLite output) then
    CoverageUtils.load_from_sqlite_database to extract coverage
    for the specific function, not the whole file.
    """
    log.info("Step 0: Computing coverage on original code...")

    git_checkout(repo_path, original_commit)
    try:
        from codeflash.languages.python.static_analysis.coverage_utils import prepare_coverage_files
        from codeflash.verification.coverage_utils import CoverageUtils

        coverage_database_file, coverage_config_file = prepare_coverage_files()

        env = os.environ.copy()
        env["PYTHONPATH"] = str(repo_path)
        env["PYTEST_PLUGINS"] = "codeflash.verification.pytest_plugin"
        env["CODEFLASH_TEST_ITERATION"] = "99"
        env["CODEFLASH_LOOP_INDEX"] = "1"
        env["CODEFLASH_TRACER_DISABLE"] = "1"

        # Run with coverage.py using codeflash's rcfile (branch coverage, SQLite output)
        cmd = [
            sys.executable, "-m", "coverage", "run",
            f"--rcfile={coverage_config_file.as_posix()}",
            "-m", "pytest",
        ]
        cmd.extend([str(f) for f in behavioral_tests])
        cmd.extend([
            "--capture=tee-sys", "-q", "--no-header",
            "--codeflash_loops_scope=session",
            "--codeflash_min_loops=1", "--codeflash_max_loops=1", "--codeflash_seconds=30",
            "-p", "no:benchmark", "-p", "no:codspeed", "-p", "no:xdist", "-p", "no:sugar",
            "--override-ini=addopts=",
        ])

        cov_run = subprocess.run(cmd, cwd=repo_path, capture_output=True, text=True, encoding="utf-8", timeout=120, env=env)
        if DEBUG:
            log.debug(f"  [coverage] run rc={cov_run.returncode}")
            log.debug(f"  [coverage] db exists={coverage_database_file.exists()}")
            if cov_run.returncode != 0:
                log.debug(f"  [coverage] stderr: {cov_run.stderr[-300:]}")

        # Load coverage from SQLite, generate JSON report, extract function-level coverage
        from coverage import Coverage
        from coverage.jsonreport import JsonReporter

        source_file = repo_path / file_path
        cov = Coverage(data_file=coverage_database_file, config_file=coverage_config_file,
                       data_suffix=True, auto_data=True, branch=True)
        cov.load()

        report_file = tmp_dir / "coverage_report.json"
        with report_file.open("w", encoding="utf-8") as f:
            try:
                JsonReporter(cov).report(morfs=[source_file.as_posix()], outfile=f)
            except Exception:
                return None

        import json as json_mod
        report = json_mod.loads(report_file.read_text(encoding="utf-8"))

        # Find the file's coverage data
        for file_key, file_data in report.get("files", {}).items():
            executed = set(file_data.get("executed_lines", []))
            missing = set(file_data.get("missing_lines", []))

            # Get function line range from AST
            src = source_file.read_text(encoding="utf-8")
            import ast as ast_mod
            func_start = func_end = None
            for node in ast_mod.walk(ast_mod.parse(src)):
                if isinstance(node, (ast_mod.FunctionDef, ast_mod.AsyncFunctionDef)) and node.name == function_name:
                    func_start = node.lineno
                    func_end = node.end_lineno
                    break

            if func_start and func_end:
                func_executed = {l for l in executed if func_start <= l <= func_end}
                func_missing = {l for l in missing if func_start <= l <= func_end}
                total = func_executed | func_missing
                if total:
                    return round(len(func_executed) / len(total) * 100, 1)
            break
    except Exception as e:
        log.warning(f"  [coverage] failed: {e}")
        if DEBUG:
            import traceback
            log.debug(traceback.format_exc())
    finally:
        git_checkout(repo_path, candidate_commit)

    return None


# ---------------------------------------------------------------------------
# Main evaluation flow
# ---------------------------------------------------------------------------


def evaluate(
    function_name: str,
    file_path: str,
    repo_path: Path,
    test_dir: Path,
    output_dir: Path,
    coverage_pct: float | None = None,
) -> None:
    target_file = repo_path / file_path
    fail_reward = dict(correct=0.0, speedup=0.0, speedup_per_test={}, tests_passed=0, tests_total=0)

    log.info(f"Evaluating: function={function_name} file={file_path}")
    log.info(f"  repo={repo_path} test_dir={test_dir} output={output_dir}")
    if DEBUG:
        log.debug(f"  python: {sys.executable}")
        log.debug(f"  target file exists: {target_file.exists()}")

    if not target_file.exists():
        log.error(f"Target file not found: {target_file}")
        write_reward(output_dir, **fail_reward)
        return

    # Git-based code switching:
    # The agent has committed on top of the pre-optimization commit.
    # HEAD = candidate (agent's optimization), HEAD~1 = original (pre-optimization)
    candidate_commit = git_get_head(repo_path)
    original_commit = git_get_parent(repo_path, candidate_commit)

    if DEBUG:
        log.debug(f"  candidate commit: {candidate_commit[:12]}")
        log.debug(f"  original commit:  {original_commit[:12]}")

    # Read candidate function code for quality review later
    candidate_file_content = target_file.read_text(encoding="utf-8")
    # Read original function code from original.py (for quality review comparison)
    original_func_file = test_dir / "original.py"
    original_func_code = original_func_file.read_text(encoding="utf-8") if original_func_file.exists() else ""

    if DEBUG:
        info = find_function_in_source(candidate_file_content, function_name)
        if info:
            start, end, _ = info
            cand_lines = candidate_file_content.splitlines()
            log.debug(f"  candidate function at lines {start + 1}-{end}:")
            log.debug("\n".join(cand_lines[start : min(end, start + 10)]))
        else:
            log.debug(f"  WARNING: function '{function_name}' not found in candidate file!")

    with tempfile.TemporaryDirectory(prefix="codeflash_eval_") as tmp_str:
        tmp_dir = Path(tmp_str)

        if DEBUG:
            log.debug(f"  tmp_dir: {tmp_dir}")

        # Discover test files from source dir (not yet placed)
        behavioral_src, perf_src = discover_test_files(test_dir)

        if not behavioral_src:
            log.error("No behavioral test files found")
            write_reward(output_dir, **fail_reward)
            return

        log.info(f"Found {len(behavioral_src)} behavioral, {len(perf_src)} perf test files")

        # Place behavioral tests for the behavioral phase
        behavioral_tests = place_test_files(test_dir, repo_path, "*behaviorinstrumented*.py")

        # Build TestFiles/TestConfig for codeflash's parser
        cf_test_files, cf_test_config = build_test_files_and_config(
            behavioral_tests, behavioral_tests, repo_path,
        )

        # ---- Step 0: Compute coverage on original code (if not precomputed) ----
        if coverage_pct is None:
            coverage_pct = compute_coverage(
                behavioral_tests, repo_path, file_path, function_name, tmp_dir, original_commit, candidate_commit,
            )
            if coverage_pct is not None:
                log.info(f"  Coverage: {coverage_pct:.1f}%")

        # ---- Step 1: Behavioral tests on ORIGINAL code ----
        log.info("Step 1: Behavioral tests on original code...")
        git_checkout(repo_path, original_commit)
        orig_proc, orig_xml, _ = run_tests(behavioral_tests, repo_path, tmp_dir, iteration=0, mode="behavior")
        git_checkout(repo_path, candidate_commit)

        orig_results = run_parse_test_results(
            orig_xml, cf_test_files, cf_test_config, iteration=0,
            function_name=function_name, source_file=target_file, run_result=orig_proc,
        )
        if len(orig_results) == 0:
            log.error("  Original code produced no test results")
            log.error(f"  stdout (last 1000): {orig_proc.stdout[-1000:]}")
            log.error(f"  stderr (last 500): {orig_proc.stderr[-500:]}")
            write_reward(output_dir, **fail_reward)
            return
        log.info(f"  Original: {len(orig_results)} invocations, {orig_results.number_of_loops()} loops")

        # ---- Step 2: Behavioral tests on CANDIDATE code ----
        log.info("Step 2: Behavioral tests on candidate code...")
        cand_proc, cand_xml, _ = run_tests(behavioral_tests, repo_path, tmp_dir, iteration=1, mode="behavior")
        cand_results = run_parse_test_results(
            cand_xml, cf_test_files, cf_test_config, iteration=1,
            function_name=function_name, source_file=target_file, run_result=cand_proc,
        )
        log.info(f"  Candidate: {len(cand_results)} invocations, {cand_results.number_of_loops()} loops")

        # ---- Step 3: Compare using codeflash's equivalence ----
        log.info("Step 3: Comparing (return values + mutations + exceptions + stdout)...")
        all_correct, test_diffs = compare_test_results(orig_results, cand_results)

        total = len(orig_results.get_all_unique_invocation_loop_ids())
        passed = total - len(test_diffs)

        if not all_correct:
            log.info(f"  INCORRECT: {passed}/{total} passed, {len(test_diffs)} diffs")
            for diff in test_diffs[:10]:
                log.info(f"    [{diff.scope.value}] {diff.original_value[:80]}  vs  {diff.candidate_value[:80]}")
            if len(test_diffs) > 10:
                log.info(f"    ... and {len(test_diffs) - 10} more diffs")
            write_reward(
                output_dir, correct=0.0, speedup=0.0, speedup_per_test={}, tests_passed=passed, tests_total=total
            )
            return

        log.info(f"  CORRECT: {passed}/{total} passed")

        # ---- Step 4: Performance benchmarks ----
        # Place perf test files (overwrites behavioral files at same dest dir)
        perf_tests = place_test_files(test_dir, repo_path, "*perfinstrumented*.py")
        test_files_for_perf = perf_tests if perf_tests else behavioral_tests
        label = "perf" if perf_tests else "behavioral (as perf)"
        log.info(f"Step 4: Performance benchmarks using {label} ({len(test_files_for_perf)} files)...")

        perf_test_files, _ = build_test_files_and_config(test_files_for_perf, test_files_for_perf, repo_path)

        log.info("  4a: Benchmarking original...")
        git_checkout(repo_path, original_commit)
        orig_perf_proc, orig_perf_xml, _ = run_tests(
            test_files_for_perf, repo_path, tmp_dir, iteration=10, mode="performance"
        )
        git_checkout(repo_path, candidate_commit)
        orig_perf_results = run_parse_test_results(
            orig_perf_xml, perf_test_files, cf_test_config, iteration=10,
            function_name=function_name, source_file=target_file, run_result=orig_perf_proc,
        )
        original_runtime = orig_perf_results.total_passed_runtime()
        log.info(f"  Original runtime: {original_runtime:,}ns ({original_runtime / 1e6:.2f}ms)")

        log.info("  4b: Benchmarking candidate...")
        cand_perf_proc, cand_perf_xml, _ = run_tests(
            test_files_for_perf, repo_path, tmp_dir, iteration=11, mode="performance"
        )
        cand_perf_results = run_parse_test_results(
            cand_perf_xml, perf_test_files, cf_test_config, iteration=11,
            function_name=function_name, source_file=target_file, run_result=cand_perf_proc,
        )
        candidate_runtime = cand_perf_results.total_passed_runtime()
        log.info(f"  Candidate runtime: {candidate_runtime:,}ns ({candidate_runtime / 1e6:.2f}ms)")

        if original_runtime > 0 and candidate_runtime > 0:
            overall_speedup = round(original_runtime / candidate_runtime, 4)
            per_test_speedup = compute_per_test_speedup(orig_perf_results, cand_perf_results)
        else:
            overall_speedup = 1.0
            per_test_speedup = {}
            log.warning("  Could not measure performance (zero runtime)")
            if DEBUG:
                log.debug(f"  [perf] orig invocations: {len(orig_perf_results)}, cand invocations: {len(cand_perf_results)}")
        log.info(f"  Overall speedup: {overall_speedup}x")

        # ---- Step 5: Quality rating ----
        # Extract just the function from both files for focused quality review
        orig_func = _extract_function(original_func_code, function_name) or original_func_code[:5000]
        cand_func = _extract_function(candidate_file_content, function_name) or candidate_file_content[:5000]
        quality = get_optimization_quality(
            orig_func, cand_func, file_path, overall_speedup,
            original_runtime_ns=original_runtime, candidate_runtime_ns=candidate_runtime,
        )
        if quality:
            log.info(f"  Quality: {quality}")

        # ---- Step 6: Write reward ----
        write_reward(
            output_dir,
            correct=1.0,
            speedup=overall_speedup,
            speedup_per_test=per_test_speedup,
            tests_passed=passed,
            tests_total=total,
            optimization_quality=quality,
            coverage_pct=coverage_pct,
        )


def main() -> None:
    parser = argparse.ArgumentParser(description="Codeflash RL Environment Evaluator")
    parser.add_argument("--function", required=True, help="Target function name (short name)")
    parser.add_argument("--file-path", required=True, help="Target file path relative to repo root")
    parser.add_argument("--repo-path", type=Path, default=Path("/workspace/inference"))
    parser.add_argument("--test-dir", type=Path, default=Path("/tests/codeflash_eval"))
    parser.add_argument("--output-dir", type=Path, default=Path("/logs/verifier"))
    parser.add_argument("--coverage-pct", type=float, default=None, help="Precomputed coverage percentage")
    parser.add_argument("--debug", action="store_true", help="Enable verbose debug logging")
    args = parser.parse_args()

    global DEBUG
    DEBUG = args.debug

    level = logging.DEBUG if DEBUG else logging.INFO
    logging.basicConfig(level=level, format="%(levelname)s: %(message)s", force=True)

    evaluate(
        function_name=args.function,
        file_path=args.file_path,
        repo_path=args.repo_path,
        test_dir=args.test_dir,
        output_dir=args.output_dir,
        coverage_pct=args.coverage_pct,
    )


if __name__ == "__main__":
    main()
