#!/usr/bin/env python3
"""Stage and deploy the repo as a Hugging Face Docker Space."""

from __future__ import annotations

import argparse
import shutil
import tempfile
from pathlib import Path
from typing import Iterable


ROOT = Path(__file__).resolve().parent.parent
SPACE_ASSETS = ROOT / "hf_space"

INCLUDE_PATHS = [
    "training",
    "serving",
    "scripts",
    "docs",
    "neuro_causal_genome_pipeline",
    "prediction_engine",
    "docker",
    "pyproject.toml",
    "README.md",
    "Makefile",
]

IGNORE_DIRS = {
    ".git",
    ".venv",
    ".pytest_cache",
    "__pycache__",
    ".mypy_cache",
    ".ruff_cache",
    ".local_backup",
    "verification_runs",
    "ablation_results",
}

IGNORE_SUFFIXES = {
    ".pyc",
    ".pyo",
}


def _should_ignore(path: Path) -> bool:
    parts = set(path.parts)
    if parts & IGNORE_DIRS:
        return True
    return path.suffix in IGNORE_SUFFIXES


def _copy_tree(src: Path, dest: Path) -> None:
    for item in src.rglob("*"):
        rel = item.relative_to(src)
        if _should_ignore(rel):
            continue
        target = dest / rel
        if item.is_dir():
            target.mkdir(parents=True, exist_ok=True)
            continue
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(item, target)


def stage_space(stage_dir: Path) -> None:
    for relative in INCLUDE_PATHS:
        source = ROOT / relative
        target = stage_dir / relative
        if source.is_dir():
            _copy_tree(source, target)
        elif source.is_file():
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(source, target)

    shutil.copy2(SPACE_ASSETS / "Dockerfile", stage_dir / "Dockerfile")
    shutil.copy2(SPACE_ASSETS / "README.md", stage_dir / "README.md")
    shutil.copy2(SPACE_ASSETS / "start_space.sh", stage_dir / "hf_space" / "start_space.sh")


def deploy(space_id: str, token: str | None, private: bool, dry_run: bool) -> None:
    from huggingface_hub import HfApi, create_repo

    api = HfApi(token=token)
    with tempfile.TemporaryDirectory(prefix="ane-hf-space-") as tmp:
        stage_dir = Path(tmp)
        (stage_dir / "hf_space").mkdir(parents=True, exist_ok=True)
        stage_space(stage_dir)
        if dry_run:
            print(f"[dry-run] staged Space at {stage_dir}")
            return

        create_repo(
            repo_id=space_id,
            repo_type="space",
            exist_ok=True,
            private=private,
            space_sdk="docker",
            token=token,
        )
        api.upload_folder(
            repo_id=space_id,
            repo_type="space",
            folder_path=str(stage_dir),
            commit_message="Deploy ANE KAN runtime Space",
        )
        print(f"Uploaded Space: https://huggingface.co/spaces/{space_id}")


def main() -> int:
    parser = argparse.ArgumentParser(description="Deploy the repo as a Hugging Face Space")
    parser.add_argument("--space-id", required=True, help="owner/space-name")
    parser.add_argument("--token", default=None, help="HF token; defaults to HF_TOKEN env if unset")
    parser.add_argument("--private", action="store_true", help="Create the Space as private")
    parser.add_argument("--dry-run", action="store_true", help="Stage files without uploading")
    args = parser.parse_args()

    deploy(
        space_id=args.space_id,
        token=args.token,
        private=args.private,
        dry_run=args.dry_run,
    )
    return 0


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