#!/usr/bin/env python3
"""Pull the PolyGuard Qwen 0.5B/1.5B evidence bundle from the artifact repo."""

from __future__ import annotations

import argparse
import os
from pathlib import Path
import shutil

from huggingface_hub import HfApi, snapshot_download


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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Pull PolyGuard submission evidence artifacts.")
    parser.add_argument("--artifact-repo-id", default="TheJackBright/polyguard-openenv-training-full-artifacts")
    return parser.parse_args()


def copy_tree(source: Path, target: Path) -> int:
    if not source.exists():
        return 0
    count = 0
    for path in source.rglob("*"):
        if not path.is_file() or path.name == ".DS_Store":
            continue
        dest = target / path.relative_to(source)
        dest.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(path, dest)
        count += 1
    return count


def main() -> None:
    args = parse_args()
    token = os.getenv("HF_TOKEN")
    api = HfApi(token=token)
    files = api.list_repo_files(repo_id=args.artifact_repo_id, repo_type="model", token=token)
    evidence_files = [item for item in files if item.startswith("submission_evidence/qwen_0_5b_1_5b/")]
    if not evidence_files:
        print("evidence_artifacts_not_uploaded_yet")
        return
    snapshot = Path(
        snapshot_download(
            repo_id=args.artifact_repo_id,
            repo_type="model",
            token=token,
            allow_patterns=["submission_evidence/qwen_0_5b_1_5b/**"],
        )
    )
    root = snapshot / "submission_evidence" / "qwen_0_5b_1_5b"
    reports = copy_tree(root / "reports", ROOT / "outputs" / "reports" / "submission_evidence" / "qwen_0_5b_1_5b")
    charts = copy_tree(root / "charts", ROOT / "outputs" / "plots" / "submission_evidence" / "qwen_0_5b_1_5b")
    docs = copy_tree(root / "docs", ROOT / "docs" / "results" / "submission_evidence_qwen_0_5b_1_5b")
    bundle = root / "qwen_0_5b_1_5b_evidence.zip"
    if bundle.exists():
        target = ROOT / "submission_bundle" / "qwen_0_5b_1_5b_evidence.zip"
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(bundle, target)
    print(f"pulled_reports={reports} pulled_charts={charts} pulled_docs={docs}")


if __name__ == "__main__":
    main()