"""Build an optimization catalog for roboflow/inference by cross-referencing DB and git history.

Produces optimization_catalog.json with entries for each reproducible optimization task.

Usage:
    python rl_env/catalog_builder.py \
        --inference-repo /path/to/inference \
        --output optimization_catalog.json
"""

from __future__ import annotations

import argparse
import json
import logging
import re
import subprocess
from dataclasses import asdict, dataclass, field
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"

REPO_ID = "befb30da-b170-4d52-b1c8-1601e9db22d5"  # roboflow/inference
REPO_NAME_IN_FEATURES = "inference"
GHA_PATH_PREFIX = "/home/runner/work/inference/inference/"


@dataclass
class CatalogEntry:
    trace_id: str
    function_name: str
    file_path: str  # relative to repo root
    speedup_x: float | None = None
    review_quality: str | None = None
    git_commit: str | None = None  # the optimization commit (post-optimization)
    pre_optimization_commit: str | None = None  # parent commit (repo state before optimization)
    has_original_code: bool = False
    has_optimized_code: bool = False
    has_generated_tests: bool = False
    has_instrumented_tests: bool = False
    has_perf_tests: bool = False
    original_runtime_ns: float | None = None
    best_optimization_id: str | None = None
    source: str = ""  # "db_events", "git", "db_features"
    data_completeness: float = 0.0  # 0-1 score of how complete the data is


@dataclass
class GitOptimization:
    commit_sha: str
    parent_sha: str
    function_name: str
    file_path: str  # relative to repo root
    speedup_pct_from_message: float | None
    commit_message: str
    files_changed: list[str] = field(default_factory=list)


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 query_optimization_events(conn: psycopg2.extensions.connection) -> list[dict]:
    """Query optimization_events for roboflow/inference with successful optimizations."""
    with conn.cursor() as cur:
        cur.execute(
            """
            SELECT e.trace_id, e.file_path, e.function_name, e.speedup_x, e.speedup_pct,
                   e.status, e.pr_url, e.created_at::date::text,
                   e.metadata::text
            FROM optimization_events e
            WHERE e.repository_id = %s
              AND e.is_optimization_found = true
              AND e.file_path IS NOT NULL
              AND e.file_path != ''
            ORDER BY e.speedup_x DESC NULLS LAST
            """,
            (REPO_ID,),
        )
        columns = [desc[0] for desc in cur.description]
        return [dict(zip(columns, row)) for row in cur.fetchall()]


def query_optimization_features(conn: psycopg2.extensions.connection) -> list[dict]:
    """Query optimization_features for inference repo with usable data."""
    with conn.cursor() as cur:
        cur.execute(
            """
            SELECT f.trace_id,
                   f.metadata->>'function_to_optimize' as function_name,
                   f.metadata->>'file_path' as raw_file_path,
                   f.metadata->>'best_optimization_id' as best_optimization_id,
                   f.review_quality,
                   f.original_runtime,
                   f.original_code IS NOT NULL as has_original_code,
                   f.optimizations_post IS NOT NULL as has_optimizations,
                   f.generated_test IS NOT NULL as has_generated_tests,
                   f.instrumented_generated_test IS NOT NULL as has_instrumented_tests,
                   f.instrumented_perf_test IS NOT NULL as has_perf_tests,
                   array_length(f.generated_test, 1) as num_gen_tests,
                   array_length(f.instrumented_generated_test, 1) as num_instr_tests,
                   array_length(f.instrumented_perf_test, 1) as num_perf_tests,
                   f.final_explanation IS NOT NULL as has_explanation,
                   f.created_at::date::text as created_at
            FROM optimization_features f
            WHERE f.repository = %s
            ORDER BY f.created_at DESC
            """,
            (REPO_NAME_IN_FEATURES,),
        )
        columns = [desc[0] for desc in cur.description]
        return [dict(zip(columns, row)) for row in cur.fetchall()]


def query_features_best_candidate(conn: psycopg2.extensions.connection, trace_id: str) -> dict | None:
    """For a given trace, find the best correct candidate with highest speedup."""
    with conn.cursor() as cur:
        cur.execute(
            """
            SELECT f.speedup_ratio, f.is_correct, f.metadata->>'best_optimization_id' as best_opt_id
            FROM optimization_features f
            WHERE f.trace_id = %s
            """,
            (trace_id,),
        )
        row = cur.fetchone()
        if not row:
            return None
        speedup_ratio, is_correct, best_opt_id = row
        if not speedup_ratio or not is_correct:
            return None

        best_speedup = -1.0
        best_id = None
        for opt_id, speedup in speedup_ratio.items():
            if speedup is None:
                continue
            correct = is_correct.get(opt_id, False)
            if correct and speedup > best_speedup:
                best_speedup = speedup
                best_id = opt_id

        if best_id is None:
            return None

        return {"optimization_id": best_id, "speedup_ratio": best_speedup}


def normalize_file_path(raw_path: str | None) -> str | None:
    """Strip GHA runner prefix to get relative repo path."""
    if not raw_path:
        return None
    if raw_path.startswith(GHA_PATH_PREFIX):
        return raw_path[len(GHA_PATH_PREFIX) :]
    if raw_path.startswith("inference/"):
        return raw_path
    return raw_path


def parse_git_optimizations(inference_repo: Path) -> list[GitOptimization]:
    """Parse merged ⚡️ commits from the inference repo."""
    result = subprocess.run(
        ["git", "log", "--all", "--format=%H|||%P|||%s", "--grep=⚡️"],
        cwd=inference_repo,
        capture_output=True,
        text=True,
        encoding="utf-8",
    )
    optimizations = []
    for line in result.stdout.strip().split("\n"):
        if not line or "|||" not in line:
            continue
        parts = line.split("|||", 2)
        if len(parts) < 3:
            continue
        commit_sha, parents, message = parts
        parent_sha = parents.split()[0] if parents else ""

        func_name = extract_function_from_message(message)
        speedup_pct = extract_speedup_from_message(message)

        diff_result = subprocess.run(
            ["git", "diff", "--name-only", f"{parent_sha}..{commit_sha}"],
            cwd=inference_repo,
            capture_output=True,
            text=True,
            encoding="utf-8",
        )
        files_changed = [f for f in diff_result.stdout.strip().split("\n") if f]
        python_files = [f for f in files_changed if f.endswith(".py")]
        file_path = python_files[0] if len(python_files) == 1 else (python_files[0] if python_files else "")

        if func_name and file_path:
            optimizations.append(
                GitOptimization(
                    commit_sha=commit_sha,
                    parent_sha=parent_sha,
                    function_name=func_name,
                    file_path=file_path,
                    speedup_pct_from_message=speedup_pct,
                    commit_message=message.strip(),
                    files_changed=files_changed,
                )
            )

    return optimizations


def extract_function_from_message(message: str) -> str | None:
    """Extract function/method name from commit message like '⚡️ Speed up method `Foo.bar` by 50%'."""
    patterns = [r"Speed up (?:method|function) `([^`]+)`", r"Optimize (\S+)"]
    for pattern in patterns:
        match = re.search(pattern, message)
        if match:
            return match.group(1)
    return None


def extract_speedup_from_message(message: str) -> float | None:
    """Extract speedup percentage from commit message like 'by 50%'."""
    match = re.search(r"by (\d+(?:\.\d+)?)%", message)
    if match:
        return float(match.group(1))
    return None


def build_catalog(
    events: list[dict], features: list[dict], git_opts: list[GitOptimization], conn: psycopg2.extensions.connection
) -> list[CatalogEntry]:
    """Cross-reference all sources to build the optimization catalog."""
    catalog: dict[str, CatalogEntry] = {}  # keyed by trace_id or function+file

    # Index features by trace_id
    features_by_trace: dict[str, dict] = {f["trace_id"]: f for f in features}

    # Index git optimizations by function name (may have dots like Class.method)
    git_by_func: dict[str, GitOptimization] = {}
    for g in git_opts:
        func_key = g.function_name.split(".")[-1]  # just method name
        git_by_func[func_key] = g
        git_by_func[g.function_name] = g  # also full qualified

    # 1. DB-first: optimization_events with file_path
    for event in events:
        trace_id = event["trace_id"]
        file_path = event["file_path"]
        func_name = event["function_name"]

        feature = features_by_trace.get(trace_id)

        # Try to find matching git commit
        func_short = func_name.split(".")[-1] if func_name else ""
        git_opt = git_by_func.get(func_name) or git_by_func.get(func_short)

        entry = CatalogEntry(
            trace_id=trace_id,
            function_name=func_name or "",
            file_path=file_path or "",
            speedup_x=event.get("speedup_x"),
            review_quality=feature["review_quality"] if feature else None,
            git_commit=git_opt.commit_sha if git_opt else None,
            pre_optimization_commit=git_opt.parent_sha if git_opt else None,
            has_original_code=bool(feature and feature.get("has_original_code")),
            has_optimized_code=bool(feature and feature.get("has_optimizations")),
            has_generated_tests=bool(feature and feature.get("has_generated_tests")),
            has_instrumented_tests=bool(feature and feature.get("has_instrumented_tests")),
            has_perf_tests=bool(feature and feature.get("has_perf_tests")),
            original_runtime_ns=feature["original_runtime"] if feature else None,
            best_optimization_id=feature.get("best_optimization_id") if feature else None,
            source="db_events",
        )

        # If we have features, find the best candidate
        if feature:
            best = query_features_best_candidate(conn, trace_id)
            if best:
                entry.best_optimization_id = best["optimization_id"]
                if entry.speedup_x is None:
                    entry.speedup_x = best["speedup_ratio"]

        entry.data_completeness = compute_completeness(entry)
        catalog[trace_id] = entry

    # 2. Git-first: add git optimizations not already in catalog
    for git_opt in git_opts:
        func_short = git_opt.function_name.split(".")[-1]
        already_covered = any(
            e.function_name and (func_short in e.function_name or e.function_name in git_opt.function_name)
            for e in catalog.values()
        )
        if already_covered:
            continue

        # Try to find matching feature by function name
        matching_feature = None
        for f in features:
            if f.get("function_name") and func_short in (f["function_name"] or ""):
                matching_feature = f
                break

        key = f"git_{git_opt.commit_sha[:12]}"
        entry = CatalogEntry(
            trace_id=matching_feature["trace_id"] if matching_feature else key,
            function_name=git_opt.function_name,
            file_path=git_opt.file_path,
            speedup_x=git_opt.speedup_pct_from_message / 100.0 if git_opt.speedup_pct_from_message else None,
            review_quality=matching_feature["review_quality"] if matching_feature else None,
            git_commit=git_opt.commit_sha,
            pre_optimization_commit=git_opt.parent_sha,
            has_original_code=True,  # git always has original via parent
            has_optimized_code=True,  # git always has optimized
            has_generated_tests=bool(matching_feature and matching_feature.get("has_generated_tests")),
            has_instrumented_tests=bool(matching_feature and matching_feature.get("has_instrumented_tests")),
            has_perf_tests=bool(matching_feature and matching_feature.get("has_perf_tests")),
            original_runtime_ns=matching_feature["original_runtime"] if matching_feature else None,
            source="git",
        )
        entry.data_completeness = compute_completeness(entry)
        catalog[entry.trace_id] = entry

    # 3. DB feature mining: features with quality ratings not already covered
    for feature in features:
        trace_id = feature["trace_id"]
        if trace_id in catalog:
            continue

        func_name = feature.get("function_name") or ""
        raw_path = feature.get("raw_file_path") or ""
        file_path = normalize_file_path(raw_path)

        if not func_name or not file_path:
            continue

        quality = feature.get("review_quality")
        if quality not in ("high", "medium"):
            continue

        # Check if there's at least one correct candidate with positive speedup
        best = query_features_best_candidate(conn, trace_id)
        if not best or best["speedup_ratio"] <= 0:
            continue

        entry = CatalogEntry(
            trace_id=trace_id,
            function_name=func_name,
            file_path=file_path,
            speedup_x=best["speedup_ratio"],
            review_quality=quality,
            has_original_code=bool(feature.get("has_original_code")),
            has_optimized_code=bool(feature.get("has_optimizations")),
            has_generated_tests=bool(feature.get("has_generated_tests")),
            has_instrumented_tests=bool(feature.get("has_instrumented_tests")),
            has_perf_tests=bool(feature.get("has_perf_tests")),
            original_runtime_ns=feature.get("original_runtime"),
            best_optimization_id=best["optimization_id"],
            source="db_features",
        )
        entry.data_completeness = compute_completeness(entry)
        catalog[trace_id] = entry

    return deduplicate_catalog(list(catalog.values()))


def deduplicate_catalog(entries: list[CatalogEntry]) -> list[CatalogEntry]:
    """Keep the best entry per (function_name, file_path) pair.

    Prefers: higher completeness > has git commit > higher speedup.
    """
    best: dict[tuple[str, str], CatalogEntry] = {}
    for e in entries:
        key = (e.function_name, e.file_path)
        if key not in best:
            best[key] = e
            continue
        existing = best[key]
        # Prefer higher completeness
        if e.data_completeness > existing.data_completeness:
            best[key] = e
        elif e.data_completeness == existing.data_completeness:
            # Prefer entries with git commit
            if e.git_commit and not existing.git_commit:
                best[key] = e
            # Prefer higher speedup
            elif (e.speedup_x or 0) > (existing.speedup_x or 0):
                best[key] = e

    return sorted(best.values(), key=lambda e: e.data_completeness, reverse=True)


def compute_completeness(entry: CatalogEntry) -> float:
    """Score 0-1 for how complete the data is for generating a Harbor task."""
    score = 0.0
    total = 7.0

    if entry.function_name:
        score += 1
    if entry.file_path:
        score += 1
    if entry.has_original_code or entry.pre_optimization_commit:
        score += 1
    if entry.has_optimized_code or entry.git_commit:
        score += 1
    if entry.has_instrumented_tests:
        score += 1
    if entry.speedup_x and entry.speedup_x > 0:
        score += 1
    if entry.review_quality:
        score += 1

    return round(score / total, 2)


def print_statistics(catalog: list[CatalogEntry]) -> None:
    """Print summary statistics about the catalog."""
    total = len(catalog)
    log.info("=" * 60)
    log.info(f"CATALOG STATISTICS: {total} total entries")
    log.info("=" * 60)

    by_source = {}
    for e in catalog:
        by_source.setdefault(e.source, []).append(e)
    for source, entries in sorted(by_source.items()):
        log.info(f"  Source '{source}': {len(entries)} entries")

    has_git = sum(1 for e in catalog if e.git_commit)
    has_original = sum(1 for e in catalog if e.has_original_code or e.pre_optimization_commit)
    has_optimized = sum(1 for e in catalog if e.has_optimized_code or e.git_commit)
    has_tests = sum(1 for e in catalog if e.has_instrumented_tests)
    has_perf = sum(1 for e in catalog if e.has_perf_tests)
    has_quality = sum(1 for e in catalog if e.review_quality)

    log.info(f"\n  With git commit: {has_git}")
    log.info(f"  With original code: {has_original}")
    log.info(f"  With optimized code: {has_optimized}")
    log.info(f"  With instrumented tests: {has_tests}")
    log.info(f"  With perf tests: {has_perf}")
    log.info(f"  With quality rating: {has_quality}")

    # Completeness distribution
    buckets = {"100%": 0, "80-99%": 0, "60-79%": 0, "40-59%": 0, "<40%": 0}
    for e in catalog:
        c = e.data_completeness
        if c >= 1.0:
            buckets["100%"] += 1
        elif c >= 0.8:
            buckets["80-99%"] += 1
        elif c >= 0.6:
            buckets["60-79%"] += 1
        elif c >= 0.4:
            buckets["40-59%"] += 1
        else:
            buckets["<40%"] += 1
    log.info("\n  Completeness distribution:")
    for bucket, count in buckets.items():
        log.info(f"    {bucket}: {count}")

    # Ready for Harbor (completeness >= 0.7 and has tests)
    ready = [e for e in catalog if e.data_completeness >= 0.7 and (e.has_instrumented_tests or e.git_commit)]
    log.info(f"\n  HARBOR-READY (completeness >= 0.7 + tests or git): {len(ready)}")

    log.info("\n  Top 20 entries by completeness:")
    for e in catalog[:20]:
        log.info(
            f"    [{e.data_completeness:.0%}] {e.function_name} "
            f"({e.file_path}) speedup={e.speedup_x} quality={e.review_quality} "
            f"source={e.source} git={'Y' if e.git_commit else 'N'}"
        )


def main() -> None:
    parser = argparse.ArgumentParser(description="Build optimization catalog for RL environment")
    parser.add_argument(
        "--inference-repo",
        type=Path,
        default=Path("/Users/saurabh/Dropbox/hacks/inference"),
        help="Path to local inference repo clone",
    )
    parser.add_argument(
        "--output", type=Path, default=Path("optimization_catalog.json"), help="Output catalog JSON file"
    )
    args = parser.parse_args()

    log.info("Connecting to database...")
    conn = get_db_connection()

    log.info("Querying optimization_events...")
    events = query_optimization_events(conn)
    log.info(f"  Found {len(events)} successful optimization events with file paths")

    log.info("Querying optimization_features...")
    features = query_optimization_features(conn)
    log.info(f"  Found {len(features)} optimization feature records")

    log.info(f"Parsing git history from {args.inference_repo}...")
    git_opts = parse_git_optimizations(args.inference_repo)
    log.info(f"  Found {len(git_opts)} git optimization commits")

    log.info("Building catalog...")
    catalog = build_catalog(events, features, git_opts, conn)

    conn.close()

    print_statistics(catalog)

    output_data = [asdict(e) for e in catalog]
    args.output.write_text(json.dumps(output_data, indent=2, default=str), encoding="utf-8")
    log.info(f"\nCatalog written to {args.output}")


if __name__ == "__main__":
    main()
