"""
Prepare ContextBench repositories and build LEANN indexes.

For each selected ContextBench instance:
  1. Clone the repo into <WORK_ROOT>/<instance_id>
  2. Checkout base_commit
  3. Build LEANN index under <WORK_ROOT>/<instance_id>/.leann/

Usage:
    cd scripts
    python prepare_repos_with_leann.py
"""

import json
import os
import shlex
import shutil
import subprocess
import time
from pathlib import Path
from typing import Optional

from datasets import load_dataset
from git import Repo

WORK_ROOT = os.environ.get("WORK_ROOT", "contextbench_work_dir_claude")
LEANN_BIN = os.environ.get("LEANN_BIN", "leann")
DATASET_NAME = os.environ.get("DATASET_NAME", "Contextbench/ContextBench")
DATASET_SPLIT = os.environ.get("DATASET_SPLIT", "train")
BENCH_FILTER = os.environ.get("BENCH_FILTER", "").strip()  # Verified | Pro | Poly | Multi
LEANN_SOURCE_EXTENSIONS = os.environ.get(
    "LEANN_SOURCE_EXTENSIONS",
    "py,go,js,jsx,ts,tsx,java,kt,kts,rs,rb,php,cs,c,cc,cpp,h,hpp,m,mm,swift,scala,sh,sql,lua,r",
)

# LEANN index build parameters — override via env vars for sweeping configs.
# ast-chunk-size is in non-whitespace CHARACTERS (not tokens). bge-base-en-v1.5
# has a 512-token limit; at ~1.2 tokens/char: 300 chars + 64 overlap ≈ 436 tokens.
LEANN_EMBEDDING_MODEL = os.environ.get(
    "LEANN_EMBEDDING_MODEL", "jinaai/jina-embeddings-v2-base-code"
)
LEANN_AST_CHUNK_SIZE = os.environ.get("LEANN_AST_CHUNK_SIZE", "600")
LEANN_AST_CHUNK_OVERLAP = os.environ.get("LEANN_AST_CHUNK_OVERLAP", "96")
# Set to "0" to disable vendor/generated exclusion (e.g. for ablation experiments).
LEANN_EXCLUDE_VENDOR = os.environ.get("LEANN_EXCLUDE_VENDOR", "1").strip() != "0"
# Set to "1" to exclude test files (e.g. for ablation experiments). Default off
# to avoid missing bugfix targets that touch test/fixture/spec files.
LEANN_EXCLUDE_TESTS = os.environ.get("LEANN_EXCLUDE_TESTS", "0").strip() != "0"
FAILED_INSTANCES_LOG = os.environ.get(
    "FAILED_INSTANCES_LOG", "prepare_repos_with_leann_failures.jsonl"
).strip()

# Fill in ContextBench instance_ids. Leave empty to prepare all tasks
# (optionally filtered by BENCH_FILTER).
SELECTED_IDS: list[str] = [
    # "SWE-Bench-Pro__python__maintenance__bugfix__19a1fba2",
    # "SWE-Bench-Pro__python__maintenance__bugfix__2464eadb",
    # "SWE-Bench-Pro__python__maintenance__bugfix__38dc8f4e",
    # "SWE-Bench-Pro__javascript__maintenance__bugfix__2bfb5681",
    # "SWE-Bench-Pro__python__maintenance__bugfix__71253eae",
    # "SWE-Bench-Pro__javascript__maintenance__bugfix__93b583ae",
    # "SWE-Bench-Pro__python__maintenance__bugfix__dcc84d4c",
    # "SWE-Bench-Pro__python__maintenance__bugfix__462b957d",
    # "SWE-Bench-Pro__python__maintenance__bugfix__9af74069",
    # "SWE-Bench-Pro__python__maintenance__bugfix__7b688a35",
    # "SWE-Bench-Pro__python__maintenance__bugfix__64fffdfa",
    # "SWE-Bench-Pro__python__maintenance__bugfix__22a1484c",
    # "SWE-Bench-Pro__go__maintenance__bugfix__1177cd53",
    # "SWE-Bench-Pro__python__maintenance__bugfix__a4287775",
    # "SWE-Bench-Pro__python__maintenance__bugfix__ba13492e",
    # "SWE-Bench-Pro__go__maintenance__bugfix__b91d5788",
    # "SWE-Bench-Pro__python__maintenance__bugfix__091dae2f",
    # "SWE-Bench-Pro__python__maintenance__bugfix__b6eff698",
    # "SWE-Bench-Pro__python__maintenance__bugfix__fcb506a5",
    # "SWE-Bench-Pro__python__maintenance__bugfix__3cfd9a02",
    # "SWE-Bench-Pro__python__maintenance__bugfix__4c132bfd",
    # "SWE-Bench-Pro__python__maintenance__bugfix__7c2efe8a",
    "SWE-Bench-Pro__go__maintenance__bugfix__40a717e5",
    "SWE-Bench-Pro__go__maintenance__bugfix__52d866b3",
    "SWE-Bench-Pro__go__maintenance__bugfix__720b4d92",
    "SWE-Bench-Pro__go__maintenance__bugfix__997c7afd",
    "SWE-Bench-Pro__javascript__maintenance__bugfix__82518720",
    "SWE-Bench-Pro__javascript__maintenance__bugfix__e31ec45c",
    "SWE-Bench-Pro__python__maintenance__bugfix__07bb383a",
    "SWE-Bench-Pro__python__maintenance__bugfix__0bac5789",
    "SWE-Bench-Pro__python__maintenance__bugfix__18d7bbbc",
    "SWE-Bench-Pro__python__maintenance__bugfix__1cf3e889",
    "SWE-Bench-Pro__python__maintenance__bugfix__20dad82b",
    "SWE-Bench-Pro__python__maintenance__bugfix__20f502e0",
    "SWE-Bench-Pro__python__maintenance__bugfix__509a20d9",
    "SWE-Bench-Pro__python__maintenance__bugfix__53ca6a30",
    "SWE-Bench-Pro__python__maintenance__bugfix__552343cd",
    "SWE-Bench-Pro__python__maintenance__bugfix__5b2cf9bb",
    "SWE-Bench-Pro__python__maintenance__bugfix__66e05eaa",
    "SWE-Bench-Pro__python__maintenance__bugfix__6ebb54dc",
    "SWE-Bench-Pro__python__maintenance__bugfix__87bfb374",
    "SWE-Bench-Pro__python__maintenance__bugfix__89932d58",
    "SWE-Bench-Pro__python__maintenance__bugfix__942d0b14",
    "SWE-Bench-Pro__python__maintenance__bugfix__983f2896",
    "SWE-Bench-Pro__python__maintenance__bugfix__a984b409",
    "SWE-Bench-Pro__python__maintenance__bugfix__aa07d0c3",
    "SWE-Bench-Pro__python__maintenance__bugfix__cf01f471",
    "SWE-Bench-Pro__python__maintenance__bugfix__d2506f10",
    "SWE-Bench-Pro__python__maintenance__bugfix__e579f2f0",
    "SWE-Bench-Pro__python__maintenance__bugfix__eafb1f0b",
    "SWE-Bench-Pro__python__maintenance__bugfix__ef8756b1",
    "SWE-Bench-Pro__python__maintenance__bugfix__f87209f8",
    "SWE-Bench-Pro__python__maintenance__bugfix__ff79bafd",
]

if os.environ.get("SELECTED_IDS"):
    SELECTED_IDS = [x.strip() for x in os.environ["SELECTED_IDS"].split(",") if x.strip()]

# Vendor and generated-code directory/file patterns to exclude from the index.
# These are third-party or machine-generated files that are never the target of
# a bugfix, so indexing them only adds noise to search results.
_VENDOR_DIR_PATTERNS = (
    "vendor/",
    "node_modules/",
    "third_party/",
    "thirdparty/",
    "externals/",
    ".cache/",
)
_GENERATED_FILE_PATTERNS = (
    "_pb.go",
    ".pb.go",
    "_gen.go",
    ".pb.cc",
    ".pb.h",
)
# Filenames like `zz_generated.deepcopy.go` end in `.go`, not `zz_generated`;
# match these as path substrings (controller-gen / k8s-style outputs).
_GENERATED_FILE_SUBSTRINGS = ("zz_generated",)

# Test file path/name patterns to exclude from the index.
_TEST_PATH_PATTERNS = (
    "/test/",
    "/tests/",
    "/__tests__/",
    "/spec/",
    "/testdata/",
    "/test_",
    "/fixtures/",
)
_TEST_FILE_PATTERNS = (
    "_test.py",
    "_test.go",
    ".test.js",
    ".test.ts",
    ".test.jsx",
    ".test.tsx",
    ".spec.js",
    ".spec.ts",
    ".spec.jsx",
    ".spec.tsx",
    "_spec.rb",
)


def _run_command(cmd: list[str], cwd: Path) -> subprocess.CompletedProcess:
    return subprocess.run(cmd, cwd=cwd, capture_output=True, text=True)


def _read_json_file(path: Path) -> Optional[dict]:
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return None


def _count_index_chunks(repo_dir: Path, instance_id: str) -> Optional[int]:
    ids_path = repo_dir / ".leann" / "indexes" / instance_id / "documents.ids.txt"
    if not ids_path.exists():
        return None
    with ids_path.open("r", encoding="utf-8") as f:
        return sum(1 for _ in f)


def _print_subprocess_output(label: str, text: str, max_lines: int = 20) -> None:
    lines = [line for line in (text or "").splitlines() if line.strip()]
    if not lines:
        return
    print(f"   📄 {label}:")
    for line in lines[:max_lines]:
        print(f"      {line}")
    if len(lines) > max_lines:
        print(f"      ... ({len(lines) - max_lines} more lines)")


def _write_failure_report(failures: list[dict]) -> Optional[Path]:
    if not failures:
        return None
    report_path = Path(FAILED_INSTANCES_LOG)
    report_path.parent.mkdir(parents=True, exist_ok=True)
    with report_path.open("w", encoding="utf-8") as f:
        for item in failures:
            f.write(json.dumps(item, ensure_ascii=True) + "\n")
    return report_path


def _load_tasks() -> list[dict]:
    print(f"📚 Loading dataset: {DATASET_NAME} ({DATASET_SPLIT})...")
    ds = load_dataset(DATASET_NAME, split=DATASET_SPLIT)
    tasks: list[dict] = list(ds)
    if BENCH_FILTER:
        tasks = [t for t in tasks if t.get("source", "") == BENCH_FILTER]

    if SELECTED_IDS:
        task_lookup = {t["instance_id"]: t for t in tasks}
        selected: list[dict] = []
        for iid in SELECTED_IDS:
            task = task_lookup.get(iid)
            if not task:
                print(f"⚠️ Instance not found in dataset/split/filter: {iid}")
                continue
            selected.append(task)
        return selected

    return tasks


def _is_pytest_style_test_py(normalized_path: str) -> bool:
    """True for pytest-style modules: basename test_*.py (e.g. test_foo.py)."""
    name = Path(normalized_path).name
    return name.startswith("test_") and name.lower().endswith(".py")


def build_leann_index(instance_id: str, repo_dir: Path) -> tuple[bool, Optional[str]]:
    print(f"   🔍 Building LEANN index for {instance_id}...")

    result = _run_command(["git", "ls-files"], cwd=repo_dir)
    if result.returncode != 0:
        error = f"Could not list git files: {result.stderr.strip()}"
        print(f"   ⚠️ {error}")
        return False, error

    tracked_files = [f for f in result.stdout.strip().split("\n") if f]
    allowed_exts = {
        f".{ext.strip().lstrip('.').lower()}"
        for ext in LEANN_SOURCE_EXTENSIONS.split(",")
        if ext.strip()
    }
    source_files = [f for f in tracked_files if Path(f).suffix.lower() in allowed_exts]

    # Exclude vendor/generated files — they are never bugfix targets and add noise.
    if LEANN_EXCLUDE_VENDOR:
        before = len(source_files)
        normalized = [f.replace("\\", "/") for f in source_files]
        source_files = [
            f
            for f, n in zip(source_files, normalized)
            if not any(pat in n for pat in _VENDOR_DIR_PATTERNS)
            and not any(n.endswith(pat) for pat in _GENERATED_FILE_PATTERNS)
            and not any(sub in n for sub in _GENERATED_FILE_SUBSTRINGS)
        ]
        excluded = before - len(source_files)
        if excluded:
            print(
                f"   🚫 Excluded {excluded} vendor/generated files ({before} → {len(source_files)})"
            )

    # Exclude test files — they are rarely bugfix targets and consistently rank
    # high in semantic search due to mirroring production code patterns.
    if LEANN_EXCLUDE_TESTS:
        before = len(source_files)
        normalized = [f.replace("\\", "/") for f in source_files]
        source_files = [
            f
            for f, n in zip(source_files, normalized)
            if not any(pat in n for pat in _TEST_PATH_PATTERNS)
            and not _is_pytest_style_test_py(n)
            and not any(n.endswith(pat) for pat in _TEST_FILE_PATTERNS)
        ]
        excluded = before - len(source_files)
        if excluded:
            print(f"   🚫 Excluded {excluded} test files ({before} → {len(source_files)})")

    if not source_files:
        error = f"No source files found for extensions: {sorted(allowed_exts)}"
        print(f"   ⚠️ {error}")
        return False, error

    # Derive --file-types from the actual extensions present after filtering,
    # so all indexed file types benefit from AST-aware chunking.
    indexed_exts = sorted(
        {Path(f).suffix.lstrip(".").lower() for f in source_files if Path(f).suffix}
    )
    file_types_arg = ",".join(indexed_exts)

    print(f"   📊 Found {len(source_files)} source files (types: {file_types_arg})")
    leann_cmd = [
        LEANN_BIN,
        "build",
        instance_id,
        "--docs",
        *source_files,
        "--embedding-mode",
        "sentence-transformers",
        "--embedding-model",
        LEANN_EMBEDDING_MODEL,
        "--backend",
        "hnsw",
        "--file-types",
        file_types_arg,
        "--force",
        "--ast-chunk-size",
        LEANN_AST_CHUNK_SIZE,
        "--ast-chunk-overlap",
        LEANN_AST_CHUNK_OVERLAP,
        "--use-ast-chunking",
        "--no-recompute",
    ]
    debug_cmd = [
        LEANN_BIN,
        "build",
        instance_id,
        "--docs",
        f"<{len(source_files)} files>",
        *leann_cmd[len(["leann", "build", instance_id, "--docs"]) + len(source_files) :],
    ]
    print(f"   🧪 LEANN command: {shlex.join(debug_cmd)}")

    try:
        proc = subprocess.run(
            leann_cmd,
            cwd=repo_dir,
            capture_output=True,
            text=True,
            timeout=1800,
            env={
                **os.environ,
                "LEANN_EMBEDDING_DEVICE": os.environ.get("LEANN_EMBEDDING_DEVICE", "mps"),
                "LEANN_BATCH_SIZE": os.environ.get("LEANN_BATCH_SIZE", "32"),
            },
        )
    except subprocess.TimeoutExpired:
        error = "LEANN build timed out"
        print(f"   ⏰ {error}")
        return False, error
    except Exception as e:
        error = f"LEANN error: {e}"
        print(f"   ❌ {error}")
        return False, error

    _print_subprocess_output("LEANN stdout", proc.stdout)
    _print_subprocess_output("LEANN stderr", proc.stderr)

    if proc.returncode != 0:
        stderr = (proc.stderr or "").strip()
        print(f"   ❌ LEANN build failed: {stderr}")
        return False, f"LEANN build failed: {stderr}"

    meta_path = repo_dir / ".leann" / "indexes" / instance_id / "documents.leann.meta.json"
    meta = _read_json_file(meta_path)
    if meta:
        print(f"   🧠 Embedding model in index: {meta.get('embedding_model', 'unknown')}")
    print("   🌲 AST chunking requested: yes (--use-ast-chunking)")
    chunk_count = _count_index_chunks(repo_dir, instance_id)
    if chunk_count is not None:
        print(f"   🧩 Indexed chunks: {chunk_count}")

    print("   ✅ LEANN index built successfully")
    return True, None


def prepare_single_task(task: dict) -> tuple[bool, Optional[str]]:
    instance_id = task["instance_id"]
    repo_url = task["repo_url"]
    base_commit = task["base_commit"]
    target_dir = Path(WORK_ROOT) / instance_id

    print(f"\n{'=' * 72}")
    print(f"📦 Preparing: {instance_id}")
    print(f"   repo: {repo_url}")
    print(f"   commit: {base_commit[:12]}...")
    print(f"{'=' * 72}")

    if not target_dir.exists():
        print(f"   📥 Cloning {repo_url}...")
        try:
            Repo.clone_from(repo_url, target_dir)
        except Exception as e:
            print(f"   ❌ Clone failed: {e}")
            return False, f"Clone failed: {e}"
    else:
        print("   ✓ Repo already exists")

    try:
        repo = Repo(target_dir)
        print(f"   🔀 Checking out {base_commit[:8]}...")
        repo.git.reset("--hard")
        repo.git.checkout(base_commit)
        repo.git.clean("-fdx", "-e", ".leann/")
        (target_dir / "PROBLEM.md").write_text(task.get("problem_statement", ""), encoding="utf-8")
    except Exception as e:
        print(f"   ❌ Checkout/clean failed: {e}")
        return False, f"Checkout/clean failed: {e}"

    return build_leann_index(instance_id, target_dir)


def main():
    print("🚀 ContextBench Repository Preparation with LEANN Indexing")
    print("=" * 72)

    leann_path = shutil.which(LEANN_BIN)
    if not leann_path:
        print("❌ LEANN not found. Install: uv tool install leann-core --with leann")
        return
    print(f"✅ LEANN found: {leann_path}")

    tasks = _load_tasks()
    if not tasks:
        print("⚠️ No tasks selected. Set SELECTED_IDS or adjust BENCH_FILTER.")
        return

    Path(WORK_ROOT).mkdir(parents=True, exist_ok=True)
    print(f"\n📂 Work root: {WORK_ROOT}")
    print(f"🎯 Tasks to prepare: {len(tasks)}")
    if BENCH_FILTER:
        print(f"🔎 Bench filter: {BENCH_FILTER}")

    success_count = 0
    fail_count = 0
    failures: list[dict] = []
    for i, task in enumerate(tasks, start=1):
        print(f"\n[{i}/{len(tasks)}]")
        succeeded, error = prepare_single_task(task)
        if succeeded:
            success_count += 1
        else:
            fail_count += 1
            failures.append(
                {
                    "instance_id": task.get("instance_id", ""),
                    "repo_url": task.get("repo_url", ""),
                    "base_commit": task.get("base_commit", ""),
                    "error": error or "unknown error",
                    "failed_at_unix": int(time.time()),
                }
            )

    print(f"\n{'=' * 72}")
    print(f"🎉 Done! ✅ {success_count} succeeded  ❌ {fail_count} failed")
    if failures:
        print("\nFailed instances:")
        for item in failures:
            print(f"  - {item['instance_id']}: {item['error']}")
        report_path = _write_failure_report(failures)
        if report_path is not None:
            print(f"\n📝 Failure report written to: {report_path}")
    print("\nNext steps:")
    print(f"  1. Verify indexes: ls {WORK_ROOT}/*/.leann")
    print("  2. Run with LEANN: LEANN_ENABLED=1 python batch_run_selected.py")
    print("  3. Baseline run:   LEANN_ENABLED=0 python batch_run_selected.py")


if __name__ == "__main__":
    main()