"""Pull instrumented test files from the database and add them to Harbor task directories.

Usage:
    python rl_env/test_puller.py \
        --resolved-catalog resolved_catalog_full.json \
        --tasks-dir codeflash-perf-tasks
"""

from __future__ import annotations

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

import psycopg2

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

DB_HOST = "codeflash-pgsql-db-prod.postgres.database.azure.com"
DB_NAME = "postgres"
DB_USER = "readonly_user"
DB_PASSWORD = "3U@L5d2B^z1+,ohC[]dP"


def get_db_connection() -> psycopg2.extensions.connection:
    return psycopg2.connect(host=DB_HOST, dbname=DB_NAME, user=DB_USER, password=DB_PASSWORD, sslmode="require")


def sanitize_task_name(function_name: str, file_path: str) -> str:
    """Must match task_generator.py's sanitize_task_name."""
    parent = Path(file_path).stem
    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]


def pull_tests_for_trace(
    conn: psycopg2.extensions.connection, trace_id: str
) -> tuple[list[str], list[str]]:
    """Pull instrumented behavioral and performance tests for a trace."""
    with conn.cursor() as cur:
        cur.execute(
            """
            SELECT instrumented_generated_test, instrumented_perf_test
            FROM optimization_features
            WHERE trace_id = %s
            """,
            (trace_id,),
        )
        row = cur.fetchone()
        if not row:
            return [], []

        behavioral = row[0] or []
        perf = row[1] or []
        return behavioral, perf


def write_test_files(
    test_dir: Path,
    behavioral_tests: list[str],
    perf_tests: list[str],
    function_name: str,
) -> tuple[int, int]:
    """Write test files to the task's test directory."""
    short_name = function_name.rsplit(".", maxsplit=1)[-1]
    written_behavioral = 0
    written_perf = 0

    for i, test_code in enumerate(behavioral_tests):
        if not test_code or not test_code.strip():
            continue
        filename = f"test_{short_name}__behaviorinstrumented_{i}.py"
        (test_dir / filename).write_text(test_code, encoding="utf-8")
        written_behavioral += 1

    from behavior_to_perf import convert_behavioral_to_perf

    # Validate all DB perf tests first — if any are malformed, discard all and derive fresh
    all_db_perf_valid = True
    for test_code in perf_tests:
        if not test_code or not test_code.strip():
            continue
        if (
            "inspect.signature" in test_code
            or "_call__bound__arguments" in test_code
            or "codeflash_duration}######!" not in test_code
            or "def codeflash_wrap(" not in test_code
        ):
            all_db_perf_valid = False
            break

    if all_db_perf_valid:
        for i, test_code in enumerate(perf_tests):
            if not test_code or not test_code.strip():
                continue
            filename = f"test_{short_name}__perfinstrumented_{i}.py"
            (test_dir / filename).write_text(test_code, encoding="utf-8")
            written_perf += 1

    # Derive perf tests from behavioral if DB had none or they were invalid
    if written_perf == 0 and written_behavioral > 0:
        for i, test_code in enumerate(behavioral_tests):
            if not test_code or not test_code.strip():
                continue
            perf_code = convert_behavioral_to_perf(test_code)
            filename = f"test_{short_name}__perfinstrumented_{i}.py"
            (test_dir / filename).write_text(perf_code, encoding="utf-8")
            written_perf += 1

    return written_behavioral, written_perf


def generate_test_sh(entry: dict, coverage_pct: float | None = None) -> str:
    """Generate test.sh that invokes the evaluator."""
    short_name = entry["function_name"].rsplit(".", maxsplit=1)[-1]
    coverage_arg = f'\\\n    --coverage-pct {coverage_pct}' if coverage_pct is not None else ""
    return f"""#!/bin/bash
set -e
cd /workspace/inference

python /tests/codeflash_eval/evaluate.py \\
    --function "{short_name}" \\
    --file-path "{entry['file_path']}" \\
    --repo-path /workspace/inference \\
    --test-dir /tests/codeflash_eval \\
    --output-dir /logs/verifier{coverage_arg}
"""


def main() -> None:
    parser = argparse.ArgumentParser(description="Pull instrumented tests from DB into task directories")
    parser.add_argument(
        "--resolved-catalog",
        type=Path,
        default=Path("resolved_catalog_full.json"),
    )
    parser.add_argument(
        "--tasks-dir",
        type=Path,
        default=Path("codeflash-perf-tasks"),
    )
    args = parser.parse_args()

    catalog = json.loads(args.resolved_catalog.read_text(encoding="utf-8"))
    log.info(f"Loaded {len(catalog)} resolved entries")

    # Load precomputed coverage if available
    coverage_map: dict[str, float] | None = None
    coverage_path = Path(__file__).parent / "coverage_map.json"
    if coverage_path.exists():
        coverage_map = json.loads(coverage_path.read_text(encoding="utf-8"))
        log.info(f"Loaded coverage map with {len(coverage_map)} entries")

    conn = get_db_connection()

    stats = {"with_behavioral": 0, "with_perf": 0, "missing": 0, "no_task_dir": 0}

    for i, entry in enumerate(catalog):
        trace_id = entry["trace_id"]
        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():
            stats["no_task_dir"] += 1
            continue

        # Pull tests from DB
        behavioral, perf = pull_tests_for_trace(conn, trace_id)

        if not behavioral and not perf:
            stats["missing"] += 1
            continue

        # Write test files
        n_behavioral, n_perf = write_test_files(test_dir, behavioral, perf, func_name)

        if n_behavioral > 0:
            stats["with_behavioral"] += 1
        if n_perf > 0:
            stats["with_perf"] += 1

        # Update test.sh to use the evaluator (include precomputed coverage if available)
        cov = coverage_map.get(trace_id) if coverage_map else None
        (task_dir / "tests" / "test.sh").write_text(generate_test_sh(entry, coverage_pct=cov), encoding="utf-8")

        # Copy the evaluator script
        evaluator_src = Path(__file__).parent / "evaluate.py"
        if evaluator_src.exists():
            (test_dir / "evaluate.py").write_text(evaluator_src.read_text(encoding="utf-8"), encoding="utf-8")

        if (i + 1) % 50 == 0:
            log.info(f"  Processed {i + 1}/{len(catalog)}...")

    conn.close()

    log.info("=" * 60)
    log.info(f"TEST PULL RESULTS:")
    log.info(f"  With behavioral tests: {stats['with_behavioral']}")
    log.info(f"  With perf tests: {stats['with_perf']}")
    log.info(f"  Missing tests in DB: {stats['missing']}")
    log.info(f"  No task directory: {stats['no_task_dir']}")


if __name__ == "__main__":
    main()
