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

import argparse
import json
import re
import subprocess
import sys
from pathlib import Path
from typing import Any


FORBIDDEN_PATTERNS: tuple[tuple[str, str, str], ...] = (
    (
        "submission_status_language",
        r"\bA\*|\bQ1\b|Best Paper",
        "Submission draft must not contain internal venue/readiness language.",
    ),
    (
        "future_work_should",
        r"future work should",
        "Main method text must not defer the method to future work.",
    ),
    (
        "generic_q_phi_generator",
        r"q_phi|q_\s*\\phi|q_\s*\{\\phi\}|q\^\s*\\phi",
        "The main generator must be CTT transport, not a generic q_phi noise model.",
    ),
    (
        "distance_proxy_called_ptr",
        r"(PTR[^.\n]{0,80}(distance|proxy))|((distance|proxy)[^.\n]{0,80}PTR)",
        "Distance-only support diagnostics must be called PPTC, not PTR.",
    ),
)


REQUIRED_PHRASES: tuple[tuple[str, str, str], ...] = (
    (
        "title_ctt",
        "Causal Tangent Transport",
        "Paper title/introduction names CTT as the central method.",
    ),
    (
        "transport_operator",
        "T_{\\phi}(z_s,z_t,\\xi_s^+)",
        "Paper defines train-positive source-to-target transport.",
    ),
    (
        "outcome_ptr",
        "OutcomePTR",
        "Paper names measured rollout positive-tangent recall separately.",
    ),
    (
        "pptc",
        "PPTC",
        "Paper names distance-only Proxy Positive Tangent Coverage.",
    ),
    (
        "support_gap",
        "SupportGap",
        "Paper exposes the support part of CAR.",
    ),
    (
        "selector_gap",
        "SelectorGap",
        "Paper exposes the selector part of CAR.",
    ),
)


REQUIRED_PATHS: tuple[tuple[str, str, str], ...] = (
    ("metrics_module", "cil/metrics.py", "Canonical measured/proxy metrics."),
    ("metrics_eval", "scripts/eval_metrics.py", "JSON/TeX metric export and proxy guards."),
    ("metrics_tests", "tests/test_metrics.py", "Regression tests for metric separation."),
    ("chart_export", "scripts/export_cil_charts.py", "Chart database export."),
    ("chart_audit", "scripts/audit_cil_charts.py", "Chart leakage audit."),
    ("data_accounting", "runs/data_accounting/table.json", "Scripted data accounting."),
    ("data_accounting_table", "runs/data_accounting/table.tex", "Data accounting table."),
    ("leakage_audit", "runs/leakage_audit/report.json", "Leakage audit artifact."),
    ("ctt_model", "cil/models/ctt.py", "Causal Tangent Transport module."),
    ("tangent_encoder", "cil/models/tangent_encoder.py", "Tangent-code encoder/decoder helpers."),
    ("chart_encoder", "cil/models/chart_encoder.py", "Chart encoder."),
    ("utility_energy", "cil/models/utility_energy.py", "Utility energy scorer."),
    ("ctt_train", "scripts/train_ctt.py", "CTT training script."),
    ("ctt_proxy_eval", "scripts/eval_ctt_proxy.py", "Proxy support evaluation."),
    ("ctt_rollout_eval", "scripts/eval_ctt_rollout.py", "Measured rollout evaluation."),
    ("utility_train", "scripts/train_utility_energy.py", "Utility energy training."),
    ("dominance_calibration", "scripts/calibrate_dominance.py", "Calibrated dominance rule."),
    ("selector_diagnostic_sweep", "scripts/build_selector_diagnostic_sweep.py", "Selector diagnostic sweep summary."),
    ("theory_tex", "paper/sections/theory.tex", "Theory section included by paper."),
    ("paper_pdf", "latex/main.pdf", "Compiled paper PDF."),
)


REQUIRED_RUN_FILES: tuple[str, ...] = (
    "table.tex",
    "metrics.json",
    "command.txt",
    "git_hash.txt",
    "data_hash.txt",
    "split_hash.txt",
)


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


INPUT_RE = re.compile(r"\\input\{([^}]+)\}")


ALLOWED_MARKDOWN_FILES: tuple[str, ...] = ("README.md",)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Audit the CTT paper against the claim-to-artifact contract. "
            "Outputs JSON and TeX so the repo can keep README.md as the only Markdown file."
        )
    )
    parser.add_argument("--repo-root", type=Path, default=Path.cwd())
    parser.add_argument("--paper", type=Path, default=Path("latex/main.tex"))
    parser.add_argument("--out-dir", type=Path, default=Path("runs/paper_ctt_audit"))
    parser.add_argument(
        "--skip-implementation-checks",
        action="store_true",
        help="Only scan paper text and input artifacts; useful for isolated unit tests.",
    )
    args = parser.parse_args(argv)

    repo_root = args.repo_root.resolve()
    paper_path = _resolve(repo_root, args.paper)
    if not paper_path.exists():
        raise SystemExit(f"paper not found: {paper_path}")
    paper_text = paper_path.read_text()

    forbidden = _forbidden_findings(paper_text)
    phrase_checks = _phrase_checks(paper_text)
    paper_inputs = _paper_input_checks(repo_root, paper_path, paper_text)
    required_paths = [] if args.skip_implementation_checks else _required_path_checks(repo_root)
    run_artifacts = _run_artifact_checks(repo_root, paper_inputs["run_dirs"])
    markdown_policy = _markdown_policy_checks(repo_root)
    summary = _summary(
        forbidden,
        phrase_checks,
        paper_inputs,
        required_paths,
        run_artifacts,
        markdown_policy,
    )

    payload = {
        "schema_version": 1,
        "audit_policy": {
            "markdown_policy": "consolidated_readme_only",
            "markdown_note": (
                "Advisor checklist asks for report.md files, but the current workspace "
                "policy keeps README.md as the only Markdown document. This audit now "
                "fails on additional Markdown files instead of regenerating report.md."
            ),
            "allowed_markdown_files": list(ALLOWED_MARKDOWN_FILES),
        },
        "paper": str(paper_path.relative_to(repo_root)),
        "summary": summary,
        "forbidden_patterns": forbidden,
        "required_phrases": phrase_checks,
        "paper_inputs": paper_inputs,
        "required_paths": required_paths,
        "run_artifacts": run_artifacts,
        "markdown_policy": markdown_policy,
    }

    out_dir = _resolve(repo_root, args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "audit.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    (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 / "config.yaml").write_text(
        "\n".join(
            [
                "paper: " + str(paper_path.relative_to(repo_root)),
                "markdown_policy: consolidated_readme_only",
                "strict_report_md: false",
                "allowed_markdown_files:",
                *[f"  - {name}" for name in ALLOWED_MARKDOWN_FILES],
            ]
        )
        + "\n"
    )
    (out_dir / "command.txt").write_text(
        "python scripts/audit_ctt_paper_artifacts.py " + " ".join(sys.argv[1:]) + "\n"
    )
    (out_dir / "git_hash.txt").write_text(_git_hash(repo_root) + "\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")

    print(json.dumps({"out_dir": str(out_dir), **summary}, indent=2, sort_keys=True))
    return 1 if summary["status"] == "fail" else 0


def _forbidden_findings(text: str) -> list[dict[str, Any]]:
    findings: list[dict[str, Any]] = []
    for name, pattern, detail in FORBIDDEN_PATTERNS:
        for match in re.finditer(pattern, text, flags=re.IGNORECASE):
            findings.append(
                {
                    "name": name,
                    "status": "fail",
                    "detail": detail,
                    "match": match.group(0),
                    "line": _line_for_offset(text, match.start()),
                }
            )
    return findings


def _phrase_checks(text: str) -> list[dict[str, Any]]:
    checks = []
    for name, phrase, detail in REQUIRED_PHRASES:
        checks.append(
            {
                "name": name,
                "status": "pass" if phrase in text else "fail",
                "detail": detail,
                "needle": phrase,
            }
        )
    return checks


def _paper_input_checks(repo_root: Path, paper_path: Path, text: str) -> dict[str, Any]:
    rows = []
    run_dirs: list[str] = []
    for match in INPUT_RE.finditer(text):
        raw = match.group(1)
        resolved = _resolve_tex_input(paper_path.parent, raw)
        status = "pass" if resolved.exists() else "fail"
        row = {
            "raw": raw,
            "line": _line_for_offset(text, match.start()),
            "resolved": _display_path(repo_root, resolved),
            "status": status,
        }
        run_dir = _run_dir_for_input(repo_root, resolved)
        if run_dir is not None:
            row["run_dir"] = _display_path(repo_root, run_dir)
            if row["run_dir"] not in run_dirs:
                run_dirs.append(row["run_dir"])
        rows.append(row)
    return {
        "num_inputs": len(rows),
        "num_missing": sum(1 for row in rows if row["status"] != "pass"),
        "rows": rows,
        "run_dirs": sorted(run_dirs),
    }


def _required_path_checks(repo_root: Path) -> list[dict[str, Any]]:
    checks = []
    for name, rel_path, detail in REQUIRED_PATHS:
        path = repo_root / rel_path
        checks.append(
            {
                "name": name,
                "path": rel_path,
                "status": "pass" if path.exists() else "fail",
                "detail": detail,
            }
        )
    ctt_configs = sorted((repo_root / "configs/ctt").glob("*.yaml"))
    checks.append(
        {
            "name": "ctt_configs",
            "path": "configs/ctt/*.yaml",
            "status": "pass" if ctt_configs else "fail",
            "detail": "CTT loss weights and variants are config-driven.",
            "files": [_display_path(repo_root, path) for path in ctt_configs],
        }
    )
    return checks


def _run_artifact_checks(repo_root: Path, run_dirs: list[str]) -> list[dict[str, Any]]:
    rows = []
    for rel_run_dir in sorted(run_dirs):
        run_dir = repo_root / rel_run_dir
        required = list(REQUIRED_RUN_FILES)
        if rel_run_dir == "runs/data_accounting":
            required = ["table.json", "table.tex"]
        missing_required = [
            name
            for name in required
            if not (run_dir / name).exists()
            and not (rel_run_dir == "runs/paper_ctt_audit" and name == "metrics.json")
        ]
        missing_advisor = [name for name in ADVISOR_RUN_FILES if not (run_dir / name).exists()]
        status = "pass" if not missing_required else "fail"
        rows.append(
            {
                "run_dir": rel_run_dir,
                "status": status,
                "missing_required": missing_required,
                "missing_advisor_contract": missing_advisor,
                "markdown_report_policy": "consolidated_readme_only",
            }
        )
    return rows


def _markdown_policy_checks(repo_root: Path) -> dict[str, Any]:
    ignored_parts = {".git", ".venv"}
    rows = []
    for path in sorted(repo_root.rglob("*.md")):
        try:
            rel = path.relative_to(repo_root)
        except ValueError:
            continue
        if any(part in ignored_parts for part in rel.parts):
            continue
        rel_text = rel.as_posix()
        rows.append(
            {
                "path": rel_text,
                "status": "pass" if rel_text in ALLOWED_MARKDOWN_FILES else "fail",
            }
        )
    unexpected = [row["path"] for row in rows if row["status"] != "pass"]
    missing_allowed = [
        name for name in ALLOWED_MARKDOWN_FILES if not (repo_root / name).exists()
    ]
    return {
        "policy": "consolidated_readme_only",
        "allowed": list(ALLOWED_MARKDOWN_FILES),
        "rows": rows,
        "unexpected_markdown": unexpected,
        "missing_allowed": missing_allowed,
        "num_markdown_files": len(rows),
        "num_unexpected_markdown": len(unexpected),
        "status": "pass" if not unexpected and not missing_allowed else "fail",
    }


def _summary(
    forbidden: list[dict[str, Any]],
    phrase_checks: list[dict[str, Any]],
    paper_inputs: dict[str, Any],
    required_paths: list[dict[str, Any]],
    run_artifacts: list[dict[str, Any]],
    markdown_policy: dict[str, Any],
) -> dict[str, Any]:
    failures = (
        len(forbidden)
        + paper_inputs["num_missing"]
        + sum(1 for item in phrase_checks if item["status"] == "fail")
        + sum(1 for item in required_paths if item["status"] == "fail")
        + sum(1 for item in run_artifacts if item["status"] == "fail")
        + (0 if markdown_policy["status"] == "pass" else 1)
    )
    warnings = sum(len(item.get("missing_advisor_contract", [])) for item in run_artifacts)
    return {
        "status": "fail" if failures else "pass",
        "num_failures": failures,
        "num_warnings": warnings,
        "num_forbidden_matches": len(forbidden),
        "num_unexpected_markdown": markdown_policy["num_unexpected_markdown"],
        "num_paper_inputs": paper_inputs["num_inputs"],
        "num_run_dirs_in_paper": len(paper_inputs["run_dirs"]),
    }


def _latex_table(payload: dict[str, Any]) -> str:
    summary = payload["summary"]
    lines = [
        "% Auto-generated by scripts/audit_ctt_paper_artifacts.py",
        "\\begin{tabular}{lrr}",
        "\\toprule",
        "Audit item & Count & Status \\\\",
        "\\midrule",
        f"Forbidden language matches & {summary['num_forbidden_matches']} & {_latex_status(summary['num_forbidden_matches'] == 0)} \\\\",
        f"Paper inputs & {summary['num_paper_inputs']} & {_latex_status(payload['paper_inputs']['num_missing'] == 0)} \\\\",
        f"Run dirs in paper & {summary['num_run_dirs_in_paper']} & {_latex_status(all(row['status'] == 'pass' for row in payload['run_artifacts']))} \\\\",
        f"Implementation paths & {len(payload['required_paths'])} & {_latex_status(all(row['status'] == 'pass' for row in payload['required_paths']))} \\\\",
        f"Unexpected Markdown files & {summary['num_unexpected_markdown']} & {_latex_status(summary['num_unexpected_markdown'] == 0)} \\\\",
        f"Advisor-contract warnings & {summary['num_warnings']} & {'pass' if summary['num_warnings'] == 0 else 'warn'} \\\\",
        "\\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 _resolve_tex_input(base_dir: Path, raw: str) -> Path:
    candidate = (base_dir / raw).resolve()
    if candidate.exists():
        return candidate
    if candidate.suffix:
        return candidate
    return candidate.with_suffix(".tex")


def _run_dir_for_input(repo_root: Path, resolved: Path) -> Path | None:
    try:
        rel = resolved.relative_to(repo_root / "runs")
    except ValueError:
        return None
    if not rel.parts:
        return None
    return repo_root / "runs" / rel.parts[0]


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


def _line_for_offset(text: str, offset: int) -> int:
    return text.count("\n", 0, offset) + 1


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_status(ok: bool) -> str:
    return "pass" if ok else "fail"


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