import os
import sys
import time
from huggingface_hub import HfApi


def main():
    token = os.environ.get("HF_TOKEN")
    repo_id = os.environ.get("HF_SPACE_ID")
    if not token or not repo_id:
        print("Missing HF_TOKEN or HF_SPACE_ID env vars", file=sys.stderr)
        sys.exit(1)

    api = HfApi(token=token)

    base_dir = os.path.dirname(__file__)
    files_to_push = [
        (os.path.join(base_dir, "requirements.txt"), "requirements.txt"),
        (os.path.join(base_dir, "flux_pipeline.py"), "flux_pipeline.py"),
        (os.path.join(base_dir, "app.py"), "app.py"),
        (os.path.join(base_dir, "Dockerfile"), "Dockerfile"),
    ]
    for local, remote in files_to_push:
        if not os.path.exists(local):
            print(f"Missing {local}", file=sys.stderr)
            sys.exit(1)
        print(f"Uploading {remote} to {repo_id}...")
        api.upload_file(
            path_or_fileobj=local,
            path_in_repo=remote,
            repo_id=repo_id,
            repo_type="space",
        )
    print("Uploaded files. Restarting Space...")
    try:
        api.restart_space(repo_id=repo_id)
        print("Restart triggered.")
    except Exception as e:
        print(f"Warning: restart_space failed: {e}")

    # Poll runtime until RUNNING
    print("Waiting for Space to enter RUNNING stage...")
    start = time.time()
    last_stage = None
    while time.time() - start < 900:  # up to 15 minutes
        try:
            rt = api.get_space_runtime(repo_id)
            stage = getattr(rt, "stage", None)
            if stage != last_stage:
                print(f"Stage: {stage}")
                last_stage = stage
            if stage in ("RUNNING", "SUCCESS"):
                print("Space is running.")
                return
        except Exception as e:
            # transient errors during rebuild
            pass
        time.sleep(5)
    print("Timeout waiting for Space to run.", file=sys.stderr)
    sys.exit(2)


if __name__ == "__main__":
    main()