#!/usr/bin/env python
from __future__ import annotations

import argparse
import json
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


NON_MARKDOWN_BACKFILLS: tuple[str, ...] = (
    "config.yaml",
    "train.log",
    "eval.log",
    "metrics_by_task.json",
    "metrics_by_seed.json",
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Backfill transparent non-Markdown run-contract artifacts for runs "
            "referenced by the CTT paper audit. Existing files are preserved."
        )
    )
    parser.add_argument("--repo-root", type=Path, default=Path.cwd())
    parser.add_argument(
        "--audit",
        type=Path,
        default=Path("runs/paper_ctt_audit/audit.json"),
        help="Audit JSON produced by scripts/audit_ctt_paper_artifacts.py.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print planned writes without creating files.",
    )
    args = parser.parse_args(argv)

    repo_root = args.repo_root.resolve()
    audit_path = _resolve(repo_root, args.audit)
    if not audit_path.exists():
        raise SystemExit(f"audit not found: {audit_path}")
    audit = json.loads(audit_path.read_text())
    timestamp = datetime.now(timezone.utc).isoformat()
    git_hash = _git_hash(repo_root)
    writes: list[dict[str, str]] = []

    for row in audit.get("run_artifacts", []):
        run_dir = repo_root / str(row.get("run_dir", ""))
        if not run_dir.exists() or not run_dir.is_dir():
            continue
        missing = set(row.get("missing_advisor_contract", []))
        for filename in NON_MARKDOWN_BACKFILLS:
            if filename in missing and not (run_dir / filename).exists():
                content = _content_for(
                    repo_root=repo_root,
                    run_dir=run_dir,
                    filename=filename,
                    timestamp=timestamp,
                    git_hash=git_hash,
                )
                writes.append({"path": _display_path(repo_root, run_dir / filename), "kind": filename})
                if not args.dry_run:
                    (run_dir / filename).write_text(content)

    payload = {
        "schema_version": 1,
        "timestamp_utc": timestamp,
        "git_hash": git_hash,
        "audit_path": _display_path(repo_root, audit_path),
        "dry_run": args.dry_run,
        "num_writes": len(writes),
        "writes": writes,
        "markdown_policy": "report.md is intentionally not backfilled; persistent prose stays in README.md.",
    }
    out_dir = repo_root / "runs" / "paper_run_artifact_backfill"
    if not args.dry_run:
        out_dir.mkdir(parents=True, exist_ok=True)
        (out_dir / "metrics.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
        (out_dir / "table.tex").write_text(_latex_table(payload) + "\n")
        (out_dir / "command.txt").write_text(
            "python scripts/backfill_paper_run_artifacts.py " + " ".join(sys.argv[1:]) + "\n"
        )
        (out_dir / "git_hash.txt").write_text(git_hash + "\n")
        (out_dir / "data_hash.txt").write_text(_first_existing_hash(repo_root, "runs/data_accounting/table.json") + "\n")
        (out_dir / "split_hash.txt").write_text(_split_hash(repo_root) + "\n")
        (out_dir / "config.yaml").write_text(
            "\n".join(
                [
                    "audit: " + _display_path(repo_root, audit_path),
                    "markdown_policy: consolidated_readme_only",
                    "backfill_report_md: false",
                    "preserve_existing: true",
                ]
            )
            + "\n"
        )
        (out_dir / "train.log").write_text("metadata-only artifact backfill; no model training was run\n")
        (out_dir / "eval.log").write_text(
            f"backfilled {len(writes)} non-Markdown run-contract files from {audit_path}\n"
        )
        (out_dir / "metrics_by_task.json").write_text("{}\n")
        (out_dir / "metrics_by_seed.json").write_text("{}\n")

    print(json.dumps(payload, indent=2, sort_keys=True))
    return 0


def _content_for(
    *,
    repo_root: Path,
    run_dir: Path,
    filename: str,
    timestamp: str,
    git_hash: str,
) -> str:
    rel_run = _display_path(repo_root, run_dir)
    if filename == "config.yaml":
        return "\n".join(
            [
                "artifact_backfill: true",
                f"run_dir: {rel_run}",
                f"timestamp_utc: {timestamp}",
                f"git_hash: {git_hash}",
                "source: scripts/backfill_paper_run_artifacts.py",
                "note: original config was not present when the CTT paper artifact audit ran",
            ]
        ) + "\n"
    if filename == "train.log":
        return (
            "Backfilled metadata log.\n"
            "No training was executed by this backfill step.\n"
            f"Run directory: {rel_run}\n"
            f"Timestamp UTC: {timestamp}\n"
            "Original train.log was absent; see command.txt and metrics.json for reproducibility evidence.\n"
        )
    if filename == "eval.log":
        return (
            "Backfilled metadata log.\n"
            "No evaluation was executed by this backfill step.\n"
            f"Run directory: {rel_run}\n"
            f"Timestamp UTC: {timestamp}\n"
            "Original eval.log was absent; see command.txt, metrics.json, and table.tex.\n"
        )
    if filename in {"metrics_by_task.json", "metrics_by_seed.json"}:
        grouping = "task" if filename == "metrics_by_task.json" else "seed"
        return json.dumps(
            {
                "_artifact_backfill": True,
                "_grouping": grouping,
                "_run_dir": rel_run,
                "_note": (
                    "Original grouped metrics were absent. This placeholder is explicit "
                    "metadata, not a computed scientific result."
                ),
            },
            indent=2,
            sort_keys=True,
        ) + "\n"
    raise ValueError(f"unsupported backfill filename: {filename}")


def _latex_table(payload: dict[str, Any]) -> str:
    counts: dict[str, int] = {}
    for row in payload["writes"]:
        counts[row["kind"]] = counts.get(row["kind"], 0) + 1
    lines = [
        "% Auto-generated by scripts/backfill_paper_run_artifacts.py",
        "\\begin{tabular}{lr}",
        "\\toprule",
        "Artifact & Backfilled count \\\\",
        "\\midrule",
    ]
    for key in sorted(counts):
        lines.append(f"{_latex_escape(key)} & {counts[key]} \\\\")
    if not counts:
        lines.append("none & 0 \\\\")
    lines.extend(["\\bottomrule", "\\end{tabular}"])
    return "\n".join(lines)


def _resolve(repo_root: Path, path: Path) -> Path:
    return path if path.is_absolute() else repo_root / path


def _display_path(repo_root: Path, path: Path) -> str:
    try:
        return str(path.relative_to(repo_root))
    except ValueError:
        return str(path)


def _git_hash(repo_root: Path) -> str:
    try:
        return subprocess.check_output(
            ["git", "rev-parse", "HEAD"],
            cwd=repo_root,
            text=True,
            stderr=subprocess.DEVNULL,
        ).strip()
    except Exception:
        return "unknown"


def _first_existing_hash(repo_root: Path, rel_path: str) -> str:
    path = repo_root / rel_path
    if not path.exists():
        return ""
    import hashlib

    return hashlib.sha256(path.read_bytes()).hexdigest()


def _split_hash(repo_root: Path) -> str:
    path = repo_root / "runs/data_accounting/table.json"
    if not path.exists():
        return ""
    try:
        payload = json.loads(path.read_text())
    except json.JSONDecodeError:
        return ""
    return str(payload.get("split_hash", ""))


def _latex_escape(value: str) -> str:
    return value.replace("_", "\\_").replace("%", "\\%").replace("&", "\\&")


if __name__ == "__main__":
    raise SystemExit(main())