#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = ["huggingface_hub>=0.20.0"]
# ///
"""
HF Space Sandbox Manager - Create and manage dev mode Spaces.

Creates sandboxes by duplicating burtenshaw/sandbox (a dev-mode ready template).

Usage:
    uv run sandbox.py create user/my-sandbox
    uv run sandbox.py create user/my-sandbox --hardware t4-small --sleep-time 3600
    uv run sandbox.py create user/my-sandbox --from other/template --private
    uv run sandbox.py status user/my-sandbox
    uv run sandbox.py start user/my-sandbox
    uv run sandbox.py stop user/my-sandbox
    uv run sandbox.py hardware user/my-sandbox t4-medium
    uv run sandbox.py wait user/my-sandbox

Environment:
    HF_TOKEN - Hugging Face API token (or use cached credentials)
"""

import argparse
import os
import sys
import time
import uuid

from huggingface_hub import HfApi, SpaceHardware

TEMPLATE_SPACE = "burtenshaw/sandbox"

HARDWARE_MAP = {
    "cpu-basic": SpaceHardware.CPU_BASIC,
    "cpu-upgrade": SpaceHardware.CPU_UPGRADE,
    "t4-small": SpaceHardware.T4_SMALL,
    "t4-medium": SpaceHardware.T4_MEDIUM,
    "a10g-small": SpaceHardware.A10G_SMALL,
    "a10g-large": SpaceHardware.A10G_LARGE,
    "a100-large": SpaceHardware.A100_LARGE,
}


def get_api():
    token = os.environ.get("HF_TOKEN")
    return HfApi(token=token)


def cmd_create(args):
    api = get_api()
    hw = args.hardware or "cpu-basic"
    source = args.source or TEMPLATE_SPACE

    space_id = f"{args.space_id.replace('/', '-')}-sandbox-{uuid.uuid4().hex}"
    print(f"Duplicating {source} to {space_id}...")
    kwargs = {
        "from_id": source,
        "to_id": space_id,
        "private": args.private,
        "hardware": hw,
    }
    if args.sleep_time:
        kwargs["sleep_time"] = args.sleep_time
    result = api.duplicate_space(**kwargs)

    print(f"Created: {result}")
    print(f"SSH: ssh user@{args.space_id.replace('/', '-')}.hf.space")


def cmd_status(args):
    api = get_api()
    runtime = api.get_space_runtime(args.space_id)
    print(f"space: {args.space_id}")
    print(f"stage: {runtime.stage}")
    print(f"hardware: {runtime.hardware}")
    if runtime.requested_hardware:
        print(f"requested: {runtime.requested_hardware}")
    print(f"ssh: user@{args.space_id.replace('/', '-')}.hf.space")
    sys.exit(0 if runtime.stage == "RUNNING" else 1)


def cmd_start(args):
    api = get_api()
    print(f"Starting {args.space_id}...")
    api.restart_space(args.space_id)
    print("Done")


def cmd_stop(args):
    api = get_api()
    print(f"Stopping {args.space_id}...")
    api.pause_space(args.space_id)
    print("Done")


def cmd_hardware(args):
    api = get_api()
    hw = HARDWARE_MAP.get(args.hardware_type.lower())
    if not hw:
        sys.exit(f"Unknown hardware: {args.hardware_type}")

    print(f"Requesting {args.hardware_type} for {args.space_id}...")
    kwargs = {"repo_id": args.space_id, "hardware": hw}
    if args.sleep_time:
        kwargs["sleep_time"] = args.sleep_time
    api.request_space_hardware(**kwargs)
    print("Done")


def cmd_wait(args):
    api = get_api()
    deadline = time.time() + args.timeout

    print(f"Waiting for {args.space_id}...")
    while time.time() < deadline:
        runtime = api.get_space_runtime(args.space_id)
        if runtime.stage == "RUNNING":
            print(f"Running (hardware: {runtime.hardware})")
            return
        if runtime.stage in ["RUNTIME_ERROR", "BUILD_ERROR"]:
            sys.exit(f"Failed: {runtime.stage}")
        print(f"  {runtime.stage}...")
        time.sleep(args.interval)
    sys.exit("Timeout")


def main():
    parser = argparse.ArgumentParser(description="HF Space Sandbox Manager")
    sub = parser.add_subparsers(dest="cmd", required=True)

    # create
    p = sub.add_parser("create", help="Create a sandbox by duplicating template")
    p.add_argument("space_id", help="Space ID (user/name)")
    p.add_argument("--from", dest="source", help=f"Source Space (default: {TEMPLATE_SPACE})")
    p.add_argument("--hardware", default="cpu-basic", help="Hardware tier")
    p.add_argument("--sleep-time", type=int, help="Auto-sleep seconds")
    p.add_argument("--private", action="store_true")
    p.set_defaults(func=cmd_create)

    # status
    p = sub.add_parser("status", help="Get Space status")
    p.add_argument("space_id")
    p.set_defaults(func=cmd_status)

    # start
    p = sub.add_parser("start", help="Start Space")
    p.add_argument("space_id")
    p.set_defaults(func=cmd_start)

    # stop
    p = sub.add_parser("stop", help="Stop Space")
    p.add_argument("space_id")
    p.set_defaults(func=cmd_stop)

    # hardware
    p = sub.add_parser("hardware", help="Request hardware")
    p.add_argument("space_id")
    p.add_argument("hardware_type", help="Hardware tier")
    p.add_argument("--sleep-time", type=int, help="Auto-sleep seconds")
    p.set_defaults(func=cmd_hardware)

    # wait
    p = sub.add_parser("wait", help="Wait for Space to be running")
    p.add_argument("space_id")
    p.add_argument("--timeout", type=int, default=300)
    p.add_argument("--interval", type=int, default=5)
    p.set_defaults(func=cmd_wait)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()