"""Deterministically select and validate the prospective Study 2 task split.

This program reads only repository history and test outcomes. It never invokes
an LLM and never reads Study 2 result directories. Candidate rejection records
are retained so the final task sample can be audited without survivorship
ambiguity.
"""

from __future__ import annotations

import argparse
from dataclasses import asdict, dataclass
import json
from pathlib import Path
import re
import subprocess
import tempfile
import time
import tomllib
from typing import Any, Iterable

from agent_harness.repair_experiment import run_test_command
from agent_harness.specs import load_task_split, load_tasks
from agent_harness.syntax_index import parse_source_file


SINCE = "2023-01-01"
MAX_SOURCE_FILES = 8
MAX_TEST_FILES = 8
MAX_SOURCE_CHANGED_LINES = 400
TEST_TIMEOUT_SECONDS = 240
EXCLUDED_SUBJECTS = re.compile(
    r"(^|\W)(revert|renovate|dependabot|format|formatting|lint-only|generated)(\W|$)",
    re.IGNORECASE,
)
HUNK = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@")


@dataclass(frozen=True, slots=True)
class RepositoryRule:
    repository_id: str
    name: str
    repository_url: str
    path: Path
    language: str
    pinned_head: str


@dataclass(slots=True)
class CandidateAudit:
    commit: str
    parent: str | None
    title: str
    status: str
    reason: str | None
    source_paths: list[str]
    test_paths: list[str]
    source_changed_lines: int
    task_id: str | None = None
    elapsed_seconds: float = 0.0
    validation: dict[str, Any] | None = None


def run_git(repository: Path, *arguments: str, binary: bool = False) -> str | bytes:
    result = subprocess.run(
        ["git", *arguments],
        cwd=repository,
        check=False,
        capture_output=True,
        timeout=180,
    )
    if result.returncode:
        detail = result.stderr.decode("utf-8", errors="replace").strip()
        raise RuntimeError(f"git {' '.join(arguments)} failed in {repository}: {detail}")
    if binary:
        return result.stdout
    return result.stdout.decode("utf-8", errors="replace")


def load_repository(root: Path, repository_id: str) -> RepositoryRule:
    matches = sorted((root / "configs" / "repositories").glob(f"{repository_id}_*.toml"))
    if len(matches) != 1:
        raise RuntimeError(f"expected one config for {repository_id}, found {len(matches)}")
    with matches[0].open("rb") as handle:
        value = tomllib.load(handle)
    path = (root / str(value["local_path"])).resolve()
    rule = RepositoryRule(
        repository_id=str(value["repository_id"]),
        name=str(value["name"]),
        repository_url=str(value["repository_url"]),
        path=path,
        language=str(value["language"]),
        pinned_head=str(value["pinned_head"]),
    )
    if rule.repository_id != repository_id or not (path / ".git").exists():
        raise RuntimeError(f"repository config/check-out mismatch for {repository_id}")
    observed = str(run_git(path, "rev-parse", "HEAD")).strip()
    if observed != rule.pinned_head:
        raise RuntimeError(
            f"{repository_id} HEAD drifted: expected {rule.pinned_head}, observed {observed}"
        )
    return rule


def semantic_title(repository: Path, commit: str) -> str:
    message = str(run_git(repository, "show", "-s", "--format=%B", commit)).strip()
    lines = [line.strip() for line in message.splitlines() if line.strip()]
    if not lines:
        return commit
    if lines[0].lower().startswith("merge branch"):
        ignored = (
            "closes ",
            "fixes ",
            "resolves ",
            "see merge request",
            "merged-by:",
            "approved-by:",
            "reviewed-by:",
            "co-authored-by:",
        )
        title = next(
            (line for line in lines[1:] if not line.lower().startswith(ignored)),
            lines[0],
        )
    else:
        title = lines[0]
    title = re.sub(r"^(feat|fix|refactor|perf)(\([^)]*\))?!?:\s*", "", title, flags=re.I)
    title = re.sub(r"\s*\(?#\d+\)?\s*$", "", title).strip()
    return title.rstrip(". ")


def candidate_commits(rule: RepositoryRule) -> tuple[str, ...]:
    output = str(
        run_git(
            rule.path,
            "log",
            "--first-parent",
            f"--since={SINCE}",
            "--format=%H",
            rule.pinned_head,
        )
    )
    return tuple(line for line in output.splitlines() if line)


def classify_paths(rule: RepositoryRule, parent: str, commit: str) -> tuple[list[str], list[str]]:
    changed = str(run_git(rule.path, "diff", "--name-only", parent, commit)).splitlines()
    if rule.language == "go":
        tests = [path for path in changed if path.endswith("_test.go")]
        sources = [
            path
            for path in changed
            if path.endswith(".go")
            and not path.endswith("_test.go")
            and not path.endswith(".gen.go")
            and "/testdata/" not in f"/{path}"
            and "/mocks/" not in f"/{path}"
        ]
    elif rule.language == "python":
        tests = [
            path
            for path in changed
            if path.startswith("tests/unit/") and path.endswith(".py")
        ]
        sources = [path for path in changed if path.startswith("gitlab/") and path.endswith(".py")]
    else:
        raise RuntimeError(f"unsupported language: {rule.language}")
    return sorted(sources), sorted(tests)


def changed_source_lines(
    repository: Path,
    parent: str,
    commit: str,
    paths: Iterable[str],
) -> int:
    selected = tuple(paths)
    if not selected:
        return 0
    output = str(run_git(repository, "diff", "--numstat", parent, commit, "--", *selected))
    total = 0
    for line in output.splitlines():
        added, deleted, _ = line.split("\t", 2)
        if not (added.isdigit() and deleted.isdigit()):
            return MAX_SOURCE_CHANGED_LINES + 1
        total += int(added) + int(deleted)
    return total


def precheck(
    rule: RepositoryRule,
    commit: str,
    excluded_commits: set[str],
) -> CandidateAudit:
    started = time.monotonic()
    title = semantic_title(rule.path, commit)
    try:
        parent = str(run_git(rule.path, "rev-parse", f"{commit}^1")).strip()
    except RuntimeError as exc:
        return CandidateAudit(commit, None, title, "rejected", str(exc), [], [], 0)
    sources, tests = classify_paths(rule, parent, commit)
    changed_lines = changed_source_lines(rule.path, parent, commit, sources)
    reason: str | None = None
    if commit in excluded_commits:
        reason = "commit already belongs to a frozen task"
    elif EXCLUDED_SUBJECTS.search(title):
        reason = "excluded dependency/format/generated/revert subject"
    elif not sources:
        reason = "no eligible production source file"
    elif not tests:
        reason = "no eligible unit-test file"
    elif len(sources) > MAX_SOURCE_FILES:
        reason = f"production file count exceeds {MAX_SOURCE_FILES}"
    elif len(tests) > MAX_TEST_FILES:
        reason = f"test file count exceeds {MAX_TEST_FILES}"
    elif changed_lines > MAX_SOURCE_CHANGED_LINES:
        reason = f"production changed lines exceed {MAX_SOURCE_CHANGED_LINES}"
    return CandidateAudit(
        commit=commit,
        parent=parent,
        title=title,
        status="rejected" if reason else "eligible",
        reason=reason,
        source_paths=sources,
        test_paths=tests,
        source_changed_lines=changed_lines,
        elapsed_seconds=time.monotonic() - started,
    )


def test_commands(rule: RepositoryRule, test_paths: Iterable[str]) -> tuple[str, ...]:
    parents = sorted({Path(path).parent.as_posix() for path in test_paths})
    if rule.language == "go":
        return tuple(
            f"go test {'.' if parent == '.' else './' + parent} -count=1" for parent in parents
        )
    return ("python -m pytest -q " + " ".join(parents),)


def extract_snapshot(
    repository: Path,
    commit: str,
    destination: Path,
    public_url: str,
) -> None:
    """Create an isolated local clone with real Git metadata at the base SHA."""

    clone = subprocess.run(
        [
            "git",
            "clone",
            "--quiet",
            "--shared",
            "--no-checkout",
            str(repository),
            str(destination),
        ],
        check=False,
        capture_output=True,
        text=True,
        timeout=180,
    )
    if clone.returncode:
        raise RuntimeError(f"local snapshot clone failed: {clone.stderr.strip()}")
    for arguments in (
        ["checkout", "--quiet", "--detach", commit],
        ["remote", "set-url", "origin", public_url],
    ):
        result = subprocess.run(
            ["git", *arguments],
            cwd=destination,
            check=False,
            capture_output=True,
            text=True,
            timeout=180,
        )
        if result.returncode:
            raise RuntimeError(
                f"snapshot git {' '.join(arguments)} failed: {result.stderr.strip()}"
            )


def apply_patch_text(tree: Path, patch: str, filename: str) -> dict[str, Any]:
    path = tree.parent / filename
    path.write_text(patch, encoding="utf-8")
    started = time.monotonic()
    result = subprocess.run(
        ["git", "apply", "--whitespace=nowarn", str(path)],
        cwd=tree,
        check=False,
        capture_output=True,
        text=True,
        timeout=120,
    )
    return {
        "returncode": result.returncode,
        "stdout": result.stdout[-4000:],
        "stderr": result.stderr[-4000:],
        "elapsed_seconds": time.monotonic() - started,
    }


def compact_test_result(value: dict[str, Any]) -> dict[str, Any]:
    return {
        **value,
        "stdout": str(value.get("stdout", ""))[-6000:],
        "stderr": str(value.get("stderr", ""))[-6000:],
    }


def validate_candidate(
    rule: RepositoryRule,
    audit: CandidateAudit,
    source_patch: str,
    test_patch: str,
) -> dict[str, Any]:
    assert audit.parent is not None
    commands = test_commands(rule, audit.test_paths)
    started = time.monotonic()
    with tempfile.TemporaryDirectory(prefix=f"study2-{rule.repository_id.lower()}-") as temporary:
        tree = Path(temporary) / "tree"
        extract_snapshot(rule.path, audit.parent, tree, rule.repository_url)
        original = [
            compact_test_result(run_test_command(tree, command, TEST_TIMEOUT_SECONDS))
            for command in commands
        ]
        if not original or any(item["returncode"] != 0 for item in original):
            return {
                "valid": False,
                "reason": "selected public tests fail at base",
                "commands": commands,
                "original": original,
                "elapsed_seconds": time.monotonic() - started,
            }
        hidden_apply = apply_patch_text(tree, test_patch, "hidden.patch")
        if hidden_apply["returncode"] != 0:
            return {
                "valid": False,
                "reason": "hidden test patch does not apply at base",
                "commands": commands,
                "original": original,
                "hidden_apply": hidden_apply,
                "elapsed_seconds": time.monotonic() - started,
            }
        hidden = [
            compact_test_result(run_test_command(tree, command, TEST_TIMEOUT_SECONDS))
            for command in commands
        ]
        if not any(item["returncode"] != 0 for item in hidden):
            return {
                "valid": False,
                "reason": "hidden tests do not expose the base bug",
                "commands": commands,
                "original": original,
                "hidden_apply": hidden_apply,
                "hidden": hidden,
                "elapsed_seconds": time.monotonic() - started,
            }
        source_apply = apply_patch_text(tree, source_patch, "source.patch")
        if source_apply["returncode"] != 0:
            return {
                "valid": False,
                "reason": "gold production patch does not apply after hidden tests",
                "commands": commands,
                "original": original,
                "hidden_apply": hidden_apply,
                "hidden": hidden,
                "source_apply": source_apply,
                "elapsed_seconds": time.monotonic() - started,
            }
        gold = [
            compact_test_result(run_test_command(tree, command, TEST_TIMEOUT_SECONDS))
            for command in commands
        ]
        valid = bool(gold) and all(item["returncode"] == 0 for item in gold)
        return {
            "valid": valid,
            "reason": None if valid else "selected tests fail after the gold patch",
            "commands": commands,
            "original": original,
            "hidden_apply": hidden_apply,
            "hidden": hidden,
            "source_apply": source_apply,
            "gold": gold,
            "elapsed_seconds": time.monotonic() - started,
        }


def changed_new_ranges(repository: Path, parent: str, commit: str, path: str) -> list[tuple[int, int]]:
    output = str(run_git(repository, "diff", "--unified=0", parent, commit, "--", path))
    ranges: list[tuple[int, int]] = []
    for line in output.splitlines():
        match = HUNK.match(line)
        if match:
            start = int(match.group(1))
            count = int(match.group(2) or "1")
            ranges.append((start, start + max(count, 1) - 1))
    return ranges


def changed_symbols(rule: RepositoryRule, audit: CandidateAudit) -> list[str]:
    assert audit.parent is not None
    symbols: list[str] = []
    for path in audit.source_paths:
        text = str(run_git(rule.path, "show", f"{audit.commit}:{path}"))
        ranges = changed_new_ranges(rule.path, audit.parent, audit.commit, path)
        matched = [
            symbol
            for symbol in parse_source_file(path, text, rule.language)
            if any(
                symbol.line_start <= changed_end and changed_start <= symbol.line_end
                for changed_start, changed_end in ranges
            )
        ]
        symbols.extend(symbol.key for symbol in matched)
        if not matched:
            symbols.append(f"{path}::<file_scope>")
    return list(dict.fromkeys(symbols))


def toml_string(value: str) -> str:
    return json.dumps(value, ensure_ascii=False)


def toml_array(values: Iterable[str]) -> str:
    return "[" + ", ".join(toml_string(value) for value in values) + "]"


def task_statement(title: str) -> str:
    if not title:
        return "Implement the behavior described by the associated public change."
    first = title[0].upper() + title[1:]
    return first.rstrip(".") + "."


def manifest_text(
    rule: RepositoryRule,
    audit: CandidateAudit,
    task_id: str,
    source_patch_name: str,
    test_patch_name: str,
    symbols: Iterable[str],
) -> str:
    assert audit.parent is not None and audit.validation is not None
    commands = tuple(audit.validation["commands"])
    return "\n".join(
        (
            "schema_version = 1",
            f"task_id = {toml_string(task_id)}",
            f"repository_url = {toml_string(rule.repository_url)}",
            f"base_commit = {toml_string(audit.parent)}",
            f"gold_commit = {toml_string(audit.commit)}",
            f"language = {toml_string(rule.language)}",
            f"statement = {toml_string(task_statement(audit.title))}",
            f"gold_patch = {toml_string('patches/' + source_patch_name)}",
            f"test_patch = {toml_string('patches/' + test_patch_name)}",
            f"gold_files = {toml_array(audit.source_paths)}",
            f"gold_symbols = {toml_array(symbols)}",
            f"fail_to_pass_tests = {toml_array(commands)}",
            f"pass_to_pass_tests = {toml_array(commands)}",
            f"difficulty = {toml_string('single_file' if len(audit.source_paths) == 1 else 'multi_file')}",
            f"provenance = {toml_string(rule.name + ' commit ' + audit.commit + '; deterministic Study 2 held-out validation.')}",
            'validation_status = "end_to_end_ready"',
            "",
        )
    )


def existing_study2_tasks(root: Path, rule: RepositoryRule) -> list[str]:
    catalog = load_tasks(root)
    if rule.repository_id == "R001":
        frozen = load_task_split(root / "tasks" / "splits" / "end_to_end_confirmatory.txt")
        return [task_id for task_id in frozen if catalog[task_id].repository_url == rule.repository_url]
    prefix = f"TASK_S2_{rule.repository_id}_"
    return sorted(task_id for task_id in catalog if task_id.startswith(prefix))


def select(root: Path, repository_id: str, target: int, write: bool) -> dict[str, Any]:
    rule = load_repository(root, repository_id)
    existing = existing_study2_tasks(root, rule)
    if len(existing) > target:
        raise RuntimeError(f"{repository_id} already has more than target={target} tasks")
    all_tasks = load_tasks(root)
    excluded_commits = {task.gold_commit for task in all_tasks.values()}
    selected: list[str] = list(existing)
    audits: list[CandidateAudit] = []
    manifest_dir = root / "tasks" / "manifests"
    patch_dir = root / "tasks" / "patches"
    validation_dir = root / "tasks" / "validation" / "study2"
    selection_dir = root / "tasks" / "selection"
    if write:
        patch_dir.mkdir(parents=True, exist_ok=True)
        validation_dir.mkdir(parents=True, exist_ok=True)
        selection_dir.mkdir(parents=True, exist_ok=True)

    next_number = len(existing) + 1
    for commit in candidate_commits(rule):
        if len(selected) >= target:
            break
        audit = precheck(rule, commit, excluded_commits)
        audits.append(audit)
        if audit.status != "eligible":
            continue
        assert audit.parent is not None
        source_patch = str(
            run_git(rule.path, "diff", "--binary", audit.parent, commit, "--", *audit.source_paths)
        )
        test_patch = str(
            run_git(rule.path, "diff", "--binary", audit.parent, commit, "--", *audit.test_paths)
        )
        validation = validate_candidate(rule, audit, source_patch, test_patch)
        audit.validation = validation
        audit.elapsed_seconds += float(validation["elapsed_seconds"])
        if not validation["valid"]:
            audit.status = "rejected"
            audit.reason = str(validation["reason"])
            print(f"REJECT {repository_id} {commit[:12]} {audit.reason}", flush=True)
            continue

        task_id = f"TASK_S2_{repository_id}_{next_number:03d}"
        source_name = f"{task_id}_source.patch"
        test_name = f"{task_id}_tests.patch"
        symbols = changed_symbols(rule, audit)
        audit.task_id = task_id
        audit.status = "selected"
        audit.reason = None
        selected.append(task_id)
        excluded_commits.add(commit)
        next_number += 1
        if write:
            (manifest_dir / f"{task_id}.toml").write_text(
                manifest_text(rule, audit, task_id, source_name, test_name, symbols),
                encoding="utf-8",
            )
            (patch_dir / source_name).write_text(source_patch, encoding="utf-8")
            (patch_dir / test_name).write_text(test_patch, encoding="utf-8")
            (validation_dir / f"{task_id}.json").write_text(
                json.dumps(
                    {
                        "schema_version": 1,
                        "task_id": task_id,
                        "repository_id": repository_id,
                        "commit": commit,
                        "parent": audit.parent,
                        "valid_end_to_end": True,
                        "validation": validation,
                    },
                    indent=2,
                    sort_keys=True,
                )
                + "\n",
                encoding="utf-8",
            )
        print(
            f"SELECT {task_id} {commit[:12]} source={len(audit.source_paths)} "
            f"tests={len(audit.test_paths)} lines={audit.source_changed_lines}",
            flush=True,
        )

    report = {
        "schema_version": 1,
        "repository": asdict(rule) | {"path": str(rule.path)},
        "selection_policy": {
            "since": SINCE,
            "target": target,
            "max_source_files": MAX_SOURCE_FILES,
            "max_test_files": MAX_TEST_FILES,
            "max_source_changed_lines": MAX_SOURCE_CHANGED_LINES,
            "test_timeout_seconds": TEST_TIMEOUT_SECONDS,
        },
        "selected_task_ids": selected,
        "selected_count": len(selected),
        "complete": len(selected) == target,
        "audits": [asdict(item) for item in audits],
    }
    if write:
        (selection_dir / f"{repository_id}_selection.json").write_text(
            json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8"
        )
    if len(selected) != target:
        raise RuntimeError(f"{repository_id} yielded {len(selected)}/{target} valid tasks")
    return report


def finalize_split(root: Path, target: int = 20) -> tuple[str, ...]:
    tasks = load_tasks(root)
    original = list(load_task_split(root / "tasks" / "splits" / "end_to_end_confirmatory.txt"))
    by_repository: dict[str, list[str]] = {
        "R001": original,
        "R002": sorted(item for item in tasks if item.startswith("TASK_S2_R002_")),
        "R003": sorted(item for item in tasks if item.startswith("TASK_S2_R003_")),
    }
    by_repository["R001"].extend(
        sorted(item for item in tasks if item.startswith("TASK_S2_R001_"))
    )
    counts = {key: len(value) for key, value in by_repository.items()}
    if counts != {"R001": target, "R002": target, "R003": target}:
        raise RuntimeError(f"cannot freeze Study 2 split with repository counts {counts}")
    split = tuple(
        item
        for repository_id in ("R001", "R002", "R003")
        for item in by_repository[repository_id]
    )
    text = (
        "# Prospective Study 2 split; frozen before all Study 2 LLM inference.\n"
        + "\n".join(split)
        + "\n"
    )
    (root / "tasks" / "splits" / "study2_confirmatory.txt").write_text(
        text, encoding="utf-8"
    )
    return split


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1])
    parser.add_argument("--repository-id", choices=("R001", "R002", "R003"))
    parser.add_argument("--target", type=int, default=20)
    parser.add_argument("--write", action="store_true")
    parser.add_argument("--finalize-split", action="store_true")
    arguments = parser.parse_args()
    root = arguments.root.resolve()
    if arguments.finalize_split:
        split = finalize_split(root, arguments.target)
        print(json.dumps({"split_count": len(split), "task_ids": split}, indent=2))
        return
    if not arguments.repository_id:
        parser.error("--repository-id is required unless --finalize-split is used")
    report = select(root, arguments.repository_id, arguments.target, arguments.write)
    print(
        json.dumps(
            {
                "repository_id": arguments.repository_id,
                "complete": report["complete"],
                "selected_count": report["selected_count"],
                "audited_candidates": len(report["audits"]),
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()