#!/usr/bin/env python3
"""
Submit a HuggingFace Job that builds the FoundationPose base image and pushes it to Docker Hub.
"""

import argparse
import os
import sys

from huggingface_hub import run_job


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Build and push the FoundationPose base image via HuggingFace Jobs"
    )
    parser.add_argument(
        "--image-name",
        default="gpue/foundationpose-base-l2",
        help="Docker Hub image name (default: gpue/foundationpose-base-l2)",
    )
    parser.add_argument(
        "--tag",
        default="latest",
        help="Docker image tag (default: latest)",
    )
    parser.add_argument(
        "--platform",
        default="linux/amd64",
        help="Target platform for docker build (default: linux/amd64)",
    )
    parser.add_argument(
        "--dockerfile",
        default="Dockerfile.base",
        help="Dockerfile path inside repo (default: Dockerfile.base)",
    )
    parser.add_argument(
        "--context",
        default=".",
        help="Docker build context path inside repo (default: .)",
    )
    parser.add_argument(
        "--target",
        default="foundationpose-base-l2",
        help="Docker build target (default: foundationpose-base-l2)",
    )
    parser.add_argument(
        "--git-repo",
        default="https://huggingface.co/spaces/gpue/foundationpose",
        help="Git repo to clone for build context (default: HF space repo)",
    )
    parser.add_argument(
        "--git-ref",
        default="main",
        help="Git ref (branch/tag/sha) to build from (default: main)",
    )
    parser.add_argument(
        "--flavor",
        default="l40sx1",
        help="HF Jobs hardware flavor (default: l40sx1)",
    )
    parser.add_argument(
        "--timeout",
        default="2h",
        help="Job timeout (default: 2h)",
    )
    parser.add_argument("--namespace", help="Organization namespace (optional)")
    parser.add_argument(
        "--hf-token",
        help="HuggingFace token (default: from HF_TOKEN or HUGGINGFACE_TOKEN env)",
    )
    parser.add_argument(
        "--docker-user",
        default=os.getenv("DOCKER_HF_USER", "gpue"),
        help="Docker Hub username (default: DOCKER_HF_USER or gpue)",
    )
    parser.add_argument(
        "--docker-token",
        help="Docker Hub token (default: from DOCKER_HF_PAT env)",
    )

    args = parser.parse_args()

    hf_token = args.hf_token or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN")
    docker_token = args.docker_token or os.getenv("DOCKER_HF_PAT")

    if not hf_token:
        print("Error: missing HF token (set HF_TOKEN or HUGGINGFACE_TOKEN)")
        sys.exit(1)
    if not docker_token:
        print("Error: missing Docker token (set DOCKER_HF_PAT or --docker-token)")
        sys.exit(1)

    env = {
        "IMAGE_NAME": args.image_name,
        "IMAGE_TAG": args.tag,
        "PLATFORM": args.platform,
        "DOCKERFILE": args.dockerfile,
        "CONTEXT": args.context,
        "TARGET": args.target,
        "GIT_REPO": args.git_repo,
        "GIT_REF": args.git_ref,
        "DOCKER_USER": args.docker_user,
    }
    secrets = {
        "HF_TOKEN": hf_token,
        "DOCKER_TOKEN": docker_token,
    }

    command = [
        "bash",
        "-lc",
        r"""
set -euo pipefail

echo "Installing git and certificates..."
if command -v microdnf >/dev/null 2>&1; then
  microdnf install -y git ca-certificates tar gzip >/dev/null
elif command -v dnf >/dev/null 2>&1; then
  dnf install -y git ca-certificates tar gzip >/dev/null
elif command -v apt-get >/dev/null 2>&1; then
  apt-get update -qq && apt-get install -y -qq git ca-certificates tar gzip >/dev/null
fi

echo "Cloning build context..."
if [ -n "${HF_TOKEN:-}" ]; then
  AUTH_REPO=$(echo "$GIT_REPO" | sed -e "s#https://#https://user:${HF_TOKEN}@#")
  git clone --depth 1 --branch "$GIT_REF" "$AUTH_REPO" /work/repo
else
  git clone --depth 1 --branch "$GIT_REF" "$GIT_REPO" /work/repo
fi

cd /work/repo

echo "Logging in to Docker Hub..."
buildah login -u "$DOCKER_USER" -p "$DOCKER_TOKEN" docker.io

export BUILDAH_ISOLATION=chroot
export STORAGE_DRIVER=vfs

PLATFORM_OS="${PLATFORM%%/*}"
PLATFORM_ARCH="${PLATFORM##*/}"

IMAGE_REF="$IMAGE_NAME:$IMAGE_TAG"

echo "Building image $IMAGE_REF (target: $TARGET)..."
buildah bud --format docker --os "$PLATFORM_OS" --arch "$PLATFORM_ARCH" -f "$DOCKERFILE" --target "$TARGET" -t "$IMAGE_REF" "$CONTEXT"

echo "Pushing image $IMAGE_REF..."
buildah push "$IMAGE_REF" "docker://$IMAGE_REF"

echo "✓ Image pushed successfully"
""",
    ]

    print("Submitting HF job for image build...")
    print(f"  Image:      {args.image_name}:{args.tag}")
    print(f"  Target:     {args.target}")
    print(f"  Repo:       {args.git_repo}")
    print(f"  Dockerfile: {args.dockerfile}")
    print(f"  Flavor:     {args.flavor}")
    print(f"  Timeout:    {args.timeout}")
    print()

    job_info = run_job(
        image="quay.io/buildah/stable:latest",
        command=command,
        env=env,
        secrets=secrets,
        flavor=args.flavor,
        timeout=args.timeout,
        namespace=args.namespace,
    )

    print("✓ Job submitted")
    print(f"  Job ID:  {job_info.id}")
    print(f"  Job URL: {job_info.url}")
    print()
    print("Monitor logs:")
    print(f"  hf jobs logs {job_info.id}")
    print("Check status:")
    print(f"  hf jobs status {job_info.id}")


if __name__ == "__main__":
    main()