#!/usr/bin/env python3
"""Create/update the private Hugging Face Space used for evidence generation."""

from __future__ import annotations

import argparse
import os
from pathlib import Path
import shutil
import sys

from huggingface_hub import HfApi


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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Deploy the PolyGuard evidence-only Hugging Face Space.")
    parser.add_argument("--repo-id", default="TheJackBright/polyguard-openenv-evidence")
    parser.add_argument("--artifact-repo-id", default="TheJackBright/polyguard-openenv-training-full-artifacts")
    parser.add_argument("--training-space-url", default="https://thejackbright-polyguard-openenv-training-full.hf.space")
    parser.add_argument("--models", default="qwen-qwen2-5-0-5b-instruct,qwen-qwen2-5-1-5b-instruct")
    parser.add_argument("--episodes", type=int, default=8)
    parser.add_argument("--hardware", default="cpu-basic")
    parser.add_argument("--sleep-time", type=int, default=3600)
    parser.add_argument("--bundle-dir", default="/tmp/polyguard-openenv-evidence-space")
    parser.add_argument("--public", action="store_true")
    parser.add_argument("--skip-upload", action="store_true")
    parser.add_argument("--bundle-only", action="store_true")
    return parser.parse_args()


def _ignore(_dir: str, names: list[str]) -> set[str]:
    ignored = {
        ".git",
        ".venv",
        "__pycache__",
        ".pytest_cache",
        ".mypy_cache",
        ".ruff_cache",
        "outputs",
        "checkpoints",
        "polyguard_rl.egg-info",
        "dist",
        "build",
    }
    return {
        name
        for name in names
        if name in ignored
        or name.endswith(".pyc")
        or name == "node_modules"
        or name == ".DS_Store"
    }


def build_bundle(bundle_dir: Path) -> None:
    if bundle_dir.exists():
        shutil.rmtree(bundle_dir)
    shutil.copytree(ROOT, bundle_dir, ignore=_ignore)
    dockerfile = bundle_dir / "Dockerfile"
    dockerfile.write_text(
        """FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime

WORKDIR /app

ENV PYTHONUNBUFFERED=1 \\
    PIP_DISABLE_PIP_VERSION_CHECK=1 \\
    TOKENIZERS_PARALLELISM=false \\
    POLYGUARD_EVIDENCE_AUTORUN=1

COPY . .

RUN python -m pip install --upgrade pip setuptools wheel \\
    && python -m pip install --no-cache-dir -r requirements.txt \\
    && python -m pip install --no-cache-dir --no-build-isolation -e .

EXPOSE 7860

CMD ["python", "-m", "app.hf_space.evidence_runner"]
""",
        encoding="utf-8",
    )
    project_readme = bundle_dir / "PROJECT_README.md"
    if (bundle_dir / "README.md").exists():
        (bundle_dir / "README.md").replace(project_readme)
    (bundle_dir / "README.md").write_text(
        """---
title: PolyGuard Evidence Runner
sdk: docker
app_port: 7860
pinned: false
---

# PolyGuard Evidence Runner

Private Docker Space for generating PolyGuard submission evidence without retraining.

The original project README is included as `PROJECT_README.md`.
""",
        encoding="utf-8",
    )


def main() -> None:
    args = parse_args()
    bundle_dir = Path(args.bundle_dir)
    build_bundle(bundle_dir)
    if args.bundle_only:
        print(f"bundle_dir={bundle_dir}")
        return

    token = os.getenv("HF_TOKEN")
    api = HfApi(token=token)
    whoami = api.whoami(token=token)
    username = str(whoami.get("name") or whoami.get("fullname") or "")
    if username and not args.repo_id.startswith(f"{username}/"):
        print(f"[deploy_evidence_space] authenticated as {username}; target={args.repo_id}")

    space_variables = [
        {"key": "POLYGUARD_EVIDENCE_MODELS", "value": args.models},
        {"key": "POLYGUARD_EVIDENCE_EPISODES", "value": str(args.episodes)},
        {"key": "POLYGUARD_EVIDENCE_AUTORUN", "value": "1"},
        {"key": "POLYGUARD_ARTIFACT_REPO_ID", "value": args.artifact_repo_id},
        {"key": "POLYGUARD_TRAINING_SPACE_URL", "value": args.training_space_url},
        {"key": "POLYGUARD_SPACE_REPO_ID", "value": args.repo_id},
    ]
    space_secrets = [{"key": "HF_TOKEN", "value": token}] if token else None

    api.create_repo(repo_id=args.artifact_repo_id, repo_type="model", private=True, exist_ok=True)
    api.create_repo(
        repo_id=args.repo_id,
        repo_type="space",
        space_sdk="docker",
        private=not args.public,
        exist_ok=True,
        space_hardware=args.hardware,
        space_sleep_time=args.sleep_time,
        space_variables=space_variables,
        space_secrets=space_secrets,
    )

    for variable in space_variables:
        api.add_space_variable(repo_id=args.repo_id, key=variable["key"], value=variable["value"])
    if token:
        api.add_space_secret(repo_id=args.repo_id, key="HF_TOKEN", value=token)
    if not args.skip_upload:
        api.upload_folder(
            repo_id=args.repo_id,
            repo_type="space",
            folder_path=str(bundle_dir),
            commit_message="Deploy PolyGuard evidence-only Space",
            ignore_patterns=[
                ".git/*",
                ".venv/*",
                "**/node_modules/*",
                "outputs/*",
                "checkpoints/*",
                "**/__pycache__/*",
                "*.pyc",
                ".DS_Store",
                "**/.DS_Store",
            ],
        )

    try:
        api.request_space_hardware(repo_id=args.repo_id, hardware=args.hardware, sleep_time=args.sleep_time)
    except Exception as exc:  # noqa: BLE001
        print(f"hardware_request_warning={exc}", file=sys.stderr)

    print(f"space_url=https://huggingface.co/spaces/{args.repo_id}")
    print(f"artifact_repo=https://huggingface.co/{args.artifact_repo_id}")
    print(f"bundle_dir={bundle_dir}")


if __name__ == "__main__":
    main()