#!/usr/bin/env python3
"""
Deploy helper for:
1) Uploading fine-tuned weights to a model repo.
2) Uploading Space code to a Streamlit Space repo.

Usage:
  HF_TOKEN=hf_xxx python3 deploy_hf.py
"""

from __future__ import annotations

import argparse
import os
import sys
from typing import Iterable

from huggingface_hub import HfApi
from huggingface_hub.errors import BadRequestError


DEFAULT_WEIGHTS_REPO = "griddev/vlm-caption-weights"
DEFAULT_SPACE_REPO = "griddev/project_02_DS"
DEFAULT_WEIGHTS_SOURCE = "../project_02"
DEFAULT_SPACE_DIR = "."
DEFAULT_SPACE_SDK = "streamlit"


def _abort(message: str) -> None:
    print(f"ERROR: {message}", file=sys.stderr)
    raise SystemExit(1)


def _ensure_exists(path: str) -> None:
    if not os.path.exists(path):
        _abort(f"Required path not found: {path}")


def _print_list(header: str, items: Iterable[str]) -> None:
    print(header)
    for item in items:
        print(f"  - {item}")


def upload_weights(api: HfApi, repo_id: str, weights_source: str) -> None:
    source_root = os.path.abspath(weights_source)
    outputs_dir = os.path.join(source_root, "outputs")
    shakespeare_txt = os.path.join(source_root, "input.txt")
    shakespeare_weights = os.path.join(source_root, "shakespeare_transformer.pt")

    _ensure_exists(outputs_dir)
    _ensure_exists(shakespeare_txt)
    _ensure_exists(shakespeare_weights)

    api.create_repo(repo_id=repo_id, repo_type="model", exist_ok=True)
    print(f"Uploading weights to model repo: {repo_id}")

    api.upload_folder(
        repo_id=repo_id,
        repo_type="model",
        folder_path=outputs_dir,
        path_in_repo="outputs",
        commit_message="Upload fine-tuned checkpoints",
    )
    api.upload_file(
        repo_id=repo_id,
        repo_type="model",
        path_or_fileobj=shakespeare_txt,
        path_in_repo="input.txt",
        commit_message="Upload input corpus for custom VLM",
    )
    api.upload_file(
        repo_id=repo_id,
        repo_type="model",
        path_or_fileobj=shakespeare_weights,
        path_in_repo="shakespeare_transformer.pt",
        commit_message="Upload Shakespeare decoder weights",
    )


def _create_space_repo(api: HfApi, repo_id: str, space_sdk: str) -> None:
    try:
        api.create_repo(
            repo_id=repo_id,
            repo_type="space",
            space_sdk=space_sdk,
            exist_ok=True,
        )
        return
    except BadRequestError as e:
        msg = str(e)
        if "Invalid option" not in msg or "sdk" not in msg:
            raise
        print(
            "Space creation rejected sdk value "
            f"'{space_sdk}'. Retrying with 'gradio' "
            "and relying on README.md front matter for streamlit."
        )
        api.create_repo(
            repo_id=repo_id,
            repo_type="space",
            space_sdk="gradio",
            exist_ok=True,
        )


def upload_space_code(api: HfApi, repo_id: str, space_dir: str, space_sdk: str) -> None:
    space_dir = os.path.abspath(space_dir)
    _ensure_exists(space_dir)

    _create_space_repo(api, repo_id=repo_id, space_sdk=space_sdk)
    print(f"Uploading Space code to: {repo_id}")

    ignore_patterns = [
        ".git/**",
        "__pycache__/**",
        "*.pyc",
        ".DS_Store",
        "venv/**",
        "weights_bundle/**",
        "outputs/**",
        "*.pt",
    ]
    _print_list("Ignoring during Space upload:", ignore_patterns)

    api.upload_folder(
        repo_id=repo_id,
        repo_type="space",
        folder_path=space_dir,
        ignore_patterns=ignore_patterns,
        commit_message="Deploy Streamlit Space app",
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Deploy weights + Space to Hugging Face")
    parser.add_argument(
        "--weights-repo",
        default=DEFAULT_WEIGHTS_REPO,
        help=f"Model repo for fine-tuned checkpoints (default: {DEFAULT_WEIGHTS_REPO})",
    )
    parser.add_argument(
        "--space-repo",
        default=DEFAULT_SPACE_REPO,
        help=f"Space repo id (default: {DEFAULT_SPACE_REPO})",
    )
    parser.add_argument(
        "--weights-source",
        default=DEFAULT_WEIGHTS_SOURCE,
        help=f"Local folder containing outputs/, input.txt, shakespeare_transformer.pt (default: {DEFAULT_WEIGHTS_SOURCE})",
    )
    parser.add_argument(
        "--space-dir",
        default=DEFAULT_SPACE_DIR,
        help=f"Local Space code folder to upload (default: {DEFAULT_SPACE_DIR})",
    )
    parser.add_argument(
        "--space-sdk",
        default=DEFAULT_SPACE_SDK,
        help=f"Space SDK to request on create (default: {DEFAULT_SPACE_SDK})",
    )
    parser.add_argument(
        "--skip-weights",
        action="store_true",
        help="Skip uploading weights repo.",
    )
    parser.add_argument(
        "--skip-space",
        action="store_true",
        help="Skip uploading Space repo.",
    )
    parser.add_argument(
        "--token",
        default=os.getenv("HF_TOKEN"),
        help="Hugging Face token. Defaults to HF_TOKEN env var.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if not args.token:
        _abort("No token found. Set HF_TOKEN or pass --token.")

    api = HfApi(token=args.token)

    if not args.skip_weights:
        upload_weights(api, repo_id=args.weights_repo, weights_source=args.weights_source)
    if not args.skip_space:
        upload_space_code(
            api,
            repo_id=args.space_repo,
            space_dir=args.space_dir,
            space_sdk=args.space_sdk,
        )

    print("Deployment finished.")
    print(f"Weights repo: https://huggingface.co/{args.weights_repo}")
    print(f"Space repo:   https://huggingface.co/spaces/{args.space_repo}")


if __name__ == "__main__":
    main()