#!/usr/bin/env python3
"""Pull remote HF training artifacts back into the local PolyGuard repo."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import shutil

from huggingface_hub import snapshot_download


ROOT = Path(__file__).resolve().parents[1]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Download PolyGuard remote training artifacts.")
    parser.add_argument("--artifact-repo-id", default="TheJackBright/polyguard-openenv-training-full-artifacts")
    parser.add_argument("--cache-dir", default="/tmp/polyguard-training-artifacts")
    parser.add_argument(
        "--training-mode",
        choices=["auto", "full", "sft-baseline"],
        default="auto",
        help="Artifact validation mode. Auto reads outputs/reports/hf_sweep_summary.json.",
    )
    return parser.parse_args()


def _copy_tree(src: Path, dst: Path) -> None:
    if src.exists():
        dst.parent.mkdir(parents=True, exist_ok=True)
        shutil.copytree(src, dst, dirs_exist_ok=True)


def _mirror_docs_results() -> None:
    docs = ROOT / "docs" / "results"
    docs.mkdir(parents=True, exist_ok=True)
    for directory in [ROOT / "outputs" / "reports", ROOT / "outputs" / "plots"]:
        if not directory.exists():
            continue
        for path in directory.rglob("*"):
            if path.is_file() and path.suffix.lower() in {".json", ".txt", ".png"}:
                target = docs / path.relative_to(directory)
                target.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy2(path, target)


def _assert_remote_training_ready(training_mode: str = "auto") -> None:
    sweep_summary_path = ROOT / "outputs" / "reports" / "hf_sweep_summary.json"
    anti_hacking_path = ROOT / "outputs" / "reports" / "anti_hacking_overfit_report.json"
    sft_path = ROOT / "outputs" / "reports" / "sft_trl_run.json"
    grpo_path = ROOT / "outputs" / "reports" / "grpo_trl_run.json"
    postsave_path = ROOT / "outputs" / "reports" / "postsave_inference.json"
    failures: list[str] = []

    def read_json(path: Path) -> dict:
        if not path.exists():
            return {}
        return json.loads(path.read_text(encoding="utf-8"))

    sweep_summary = read_json(sweep_summary_path)
    anti_hacking = read_json(anti_hacking_path)
    if sweep_summary:
        summary_mode = str(sweep_summary.get("training_mode") or "full")
        effective_mode = summary_mode if training_mode == "auto" else training_mode
        sft_only = effective_mode == "sft-baseline"
        if int(sweep_summary.get("completed_models", 0) or 0) <= 0:
            failures.append("HF sweep has no completed models")
        for row in sweep_summary.get("models", []):
            if not isinstance(row, dict) or row.get("status") != "completed":
                continue
            label = str(row.get("label") or row.get("model_id") or "model")
            if row.get("fallback_detected"):
                failures.append(f"{label} used fallback backend")
            if not row.get("reward_range_ok"):
                failures.append(f"{label} has reward range/precision failures")
            artifact_paths = row.get("artifact_paths", {})
            if not isinstance(artifact_paths, dict):
                artifact_paths = {}
            if not artifact_paths.get("sft"):
                failures.append(f"{label} missing SFT artifact path")
            if not sft_only and not artifact_paths.get("grpo"):
                failures.append(f"{label} missing GRPO artifact path")
        charts = sweep_summary.get("charts", {})
        for chart_name, rel_path in charts.items():
            if not (ROOT / str(rel_path)).exists():
                failures.append(f"missing chart {chart_name}")
        if anti_hacking.get("passed") is not True:
            failures.append("anti-hacking/overfit report did not pass")
        if failures:
            raise SystemExit("artifact_checks_failed:" + "; ".join(failures))
        return

    sft = read_json(sft_path)
    if sft.get("status") != "ok":
        failures.append("SFT status is not ok")
    if sft.get("backend") not in {"trl_unsloth", "trl_transformers"}:
        failures.append("SFT backend is not TRL")
    if not sft.get("artifact_path"):
        failures.append("SFT artifact path is empty")
    if int(sft.get("examples_used", 0) or 0) <= 0:
        failures.append("SFT examples_used is zero")

    grpo = read_json(grpo_path)
    if grpo.get("status") != "ok":
        failures.append("GRPO status is not ok")
    if not grpo.get("artifact_path"):
        failures.append("GRPO artifact path is empty")

    postsave = read_json(postsave_path)
    if postsave.get("model_source") == "fallback_policy":
        failures.append("post-save inference still uses fallback policy")

    if failures:
        raise SystemExit("artifact_checks_failed:" + "; ".join(failures))


def main() -> None:
    args = parse_args()
    snapshot = Path(
        snapshot_download(
            repo_id=args.artifact_repo_id,
            repo_type="model",
            cache_dir=args.cache_dir,
            allow_patterns=[
                "outputs/reports/**",
                "outputs/plots/**",
                "docs/results/**",
                "checkpoints/sft_adapter/**",
                "checkpoints/grpo_adapter/**",
                "checkpoints/merged/**",
                "checkpoints/sweeps/**",
            ],
        )
    )

    for rel in [
        "outputs/reports",
        "outputs/plots",
        "docs/results",
        "checkpoints/sft_adapter",
        "checkpoints/grpo_adapter",
        "checkpoints/merged",
        "checkpoints/sweeps",
    ]:
        _copy_tree(snapshot / rel, ROOT / rel)

    _mirror_docs_results()
    _assert_remote_training_ready(training_mode=args.training_mode)
    print(f"artifacts_pulled_from={args.artifact_repo_id}")


if __name__ == "__main__":
    main()