#!/usr/bin/env python3
"""
Upload video templates to HuggingFace Space using Hub API
This bypasses Git LFS limitations and uploads files directly
"""

import os
from pathlib import Path
from huggingface_hub import HfApi, login
import sys

# Configuration
SPACE_ID = "minhho/mimo-1.0"  # Your HuggingFace Space
LOCAL_TEMPLATES_DIR = "./assets/video_template"
REMOTE_PATH_PREFIX = "assets/video_template"

# Templates to upload (you can modify this list)
TEMPLATES_TO_UPLOAD = [
    "dance_indoor_1",
    "sports_basketball_gym",
    "movie_BruceLee1",
    "shorts_kungfu_desert1",
    "shorts_kungfu_match1",
    "sports_nba_dunk",
    "sports_nba_pass",
    "parkour_climbing",
    "syn_basketball_06_13",
    "syn_dancing2_00093_irish_dance",
    "syn_football_10_05",
]

# Files to upload per template (in priority order)
FILES_TO_UPLOAD = [
    "sdc.mp4",      # REQUIRED - pose skeleton
    "config.json",  # Optional - template config
    "vid.mp4",      # Optional - original video
    "bk.mp4",       # Optional - background
    "mask.mp4",     # Optional - mask
    "occ.mp4",      # Optional - occlusion (if exists)
    "bbox.npy",     # Optional - bounding box
]


def upload_templates(token=None, templates=None, dry_run=False):
    """
    Upload video templates to HuggingFace Space

    Args:
        token: HuggingFace token (optional, will prompt if not provided)
        templates: List of template names to upload (default: all in TEMPLATES_TO_UPLOAD)
        dry_run: If True, just show what would be uploaded without actually uploading
    """

    # Login to HuggingFace
    if token:
        login(token=token)
    else:
        print("🔐 Please login to HuggingFace (you'll be prompted for your token)")
        print("   Get your token from: https://huggingface.co/settings/tokens")
        login()

    # Initialize HF API
    api = HfApi()

    # Use provided templates or default list
    templates_list = templates or TEMPLATES_TO_UPLOAD

    print(f"📦 Preparing to upload {len(templates_list)} templates to Space: {SPACE_ID}")
    print(f"📁 Local directory: {LOCAL_TEMPLATES_DIR}\n")

    if dry_run:
        print("🔍 DRY RUN MODE - No files will be uploaded\n")

    # Check local templates directory
    templates_dir = Path(LOCAL_TEMPLATES_DIR)
    if not templates_dir.exists():
        print(f"❌ Error: Templates directory not found: {LOCAL_TEMPLATES_DIR}")
        print("   Please make sure you've extracted assets.zip")
        sys.exit(1)

    uploaded_count = 0
    skipped_count = 0
    error_count = 0

    # Upload each template
    for template_name in templates_list:
        template_path = templates_dir / template_name

        if not template_path.exists():
            print(f"⚠️  Template not found: {template_name} - SKIPPED")
            skipped_count += 1
            continue

        print(f"📤 Uploading template: {template_name}")

        # Check for required sdc.mp4
        sdc_file = template_path / "sdc.mp4"
        if not sdc_file.exists():
            print(f"   ❌ Missing required file: sdc.mp4 - SKIPPED")
            skipped_count += 1
            continue

        # Upload each file in the template
        template_uploaded = False
        for file_name in FILES_TO_UPLOAD:
            file_path = template_path / file_name

            if not file_path.exists():
                continue  # Skip missing optional files

            # Calculate file size
            file_size_mb = file_path.stat().st_size / (1024 * 1024)

            # Remote path in the Space
            remote_file_path = f"{REMOTE_PATH_PREFIX}/{template_name}/{file_name}"

            print(f"   📄 {file_name} ({file_size_mb:.1f} MB)", end="")

            if dry_run:
                print(" [DRY RUN]")
                continue

            try:
                # Upload file to Space
                api.upload_file(
                    path_or_fileobj=str(file_path),
                    path_in_repo=remote_file_path,
                    repo_id=SPACE_ID,
                    repo_type="space",
                    commit_message=f"Add {template_name}/{file_name}"
                )
                print(" ✅")
                template_uploaded = True

            except Exception as e:
                print(f" ❌ Error: {str(e)[:50]}")
                error_count += 1

        if template_uploaded:
            uploaded_count += 1
            print(f"   ✅ Template uploaded successfully\n")
        else:
            print(f"   ⚠️  No files uploaded for this template\n")

    # Summary
    print("=" * 60)
    print("📊 Upload Summary:")
    print(f"   ✅ Templates uploaded: {uploaded_count}")
    print(f"   ⚠️  Templates skipped: {skipped_count}")
    print(f"   ❌ Errors: {error_count}")
    print("=" * 60)

    if not dry_run and uploaded_count > 0:
        print("\n🎉 Upload complete!")
        print(f"   Visit your Space: https://huggingface.co/spaces/{SPACE_ID}")
        print("   Click '🔄 Refresh Templates' in the app to see your templates")
    elif dry_run:
        print("\n💡 To actually upload, run without --dry-run flag")

    return uploaded_count, skipped_count, error_count


def main():
    """Main function with CLI support"""
    import argparse

    parser = argparse.ArgumentParser(
        description="Upload video templates to HuggingFace Space",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Dry run to see what would be uploaded
  python upload_templates_to_hf.py --dry-run

  # Upload all templates
  python upload_templates_to_hf.py

  # Upload specific templates only
  python upload_templates_to_hf.py --templates dance_indoor_1 sports_basketball_gym

  # Use specific HF token
  python upload_templates_to_hf.py --token YOUR_HF_TOKEN
        """
    )

    parser.add_argument(
        "--token",
        type=str,
        help="HuggingFace API token (optional, will prompt if not provided)"
    )

    parser.add_argument(
        "--templates",
        nargs="+",
        help="Specific templates to upload (default: all)"
    )

    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show what would be uploaded without actually uploading"
    )

    parser.add_argument(
        "--list",
        action="store_true",
        help="List available templates and exit"
    )

    args = parser.parse_args()

    # List templates if requested
    if args.list:
        templates_dir = Path(LOCAL_TEMPLATES_DIR)
        if templates_dir.exists():
            print("📁 Available templates:")
            for template in sorted(templates_dir.iterdir()):
                if template.is_dir() and not template.name.startswith('.'):
                    sdc_exists = (template / "sdc.mp4").exists()
                    status = "✅" if sdc_exists else "❌ (missing sdc.mp4)"
                    print(f"   {template.name} {status}")
        else:
            print(f"❌ Templates directory not found: {LOCAL_TEMPLATES_DIR}")
        return

    # Upload templates
    try:
        upload_templates(
            token=args.token,
            templates=args.templates,
            dry_run=args.dry_run
        )
    except KeyboardInterrupt:
        print("\n\n⚠️  Upload cancelled by user")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ Error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()