import os
import sys
import argparse
from huggingface_hub import HfApi, login

def upload_to_hf(local_path, repo_id, repo_type="dataset", path_in_repo=None, token=None):
    if token:
        login(token=token)
    elif not os.environ.get("HF_TOKEN"):
        print("Warning: No HF_TOKEN found. Upload might fail if the repo is private or requires authentication.")
        
    api = HfApi()
    
    if not os.path.exists(local_path):
        print(f"Error: Path does not exist: {local_path}", file=sys.stderr)
        sys.exit(1)

    # Ensure repo exists
    try:
        api.create_repo(repo_id=repo_id, repo_type=repo_type, exist_ok=True)
        print(f"Verified repository: {repo_type}s/{repo_id}")
    except Exception as e:
        print(f"Note: Could not verify or create repo (it might already exist or there's a permission issue). Continuing... ({e})")

    if os.path.isfile(local_path):
        if not path_in_repo:
            path_in_repo = os.path.basename(local_path)
        print(f"Uploading file '{local_path}' to 'hf://{repo_type}s/{repo_id}/{path_in_repo}'...")
        try:
            api.upload_file(
                path_or_fileobj=local_path,
                path_in_repo=path_in_repo,
                repo_id=repo_id,
                repo_type=repo_type
            )
            print("✅ File upload complete!")
        except Exception as e:
            print(f"❌ Failed to upload file: {e}", file=sys.stderr)
            
    elif os.path.isdir(local_path):
        if path_in_repo is None:
            path_in_repo = "" # Root of repo by default
        print(f"Syncing folder '{local_path}' to 'hf://{repo_type}s/{repo_id}/{path_in_repo}'...")
        try:
            api.upload_folder(
                folder_path=local_path,
                path_in_repo=path_in_repo,
                repo_id=repo_id,
                repo_type=repo_type
            )
            print("✅ Folder sync complete!")
        except Exception as e:
            print(f"❌ Failed to sync folder: {e}", file=sys.stderr)

def main():
    parser = argparse.ArgumentParser(description="Upload any video, file, or folder to Hugging Face")
    parser.add_argument("input_path", type=str, help="Path to the local video file or directory (e.g. ./data or video.mp4)")
    parser.add_argument("--repo", type=str, default="AdhyanshVerma/YT", help="Hugging Face repo ID (default: AdhyanshVerma/YT)")
    parser.add_argument("--repo-type", type=str, default="dataset", choices=["dataset", "model", "space"], help="Type of HF repo (default: dataset)")
    parser.add_argument("--path-in-repo", type=str, default=None, help="Destination path in the repo (defaults to filename or root of repo)")
    parser.add_argument("--token", type=str, default=os.environ.get("HF_TOKEN"), help="HF Access Token (or set HF_TOKEN env var)")
    
    args = parser.parse_args()
    
    upload_to_hf(
        local_path=args.input_path,
        repo_id=args.repo,
        repo_type=args.repo_type,
        path_in_repo=args.path_in_repo,
        token=args.token
    )

if __name__ == "__main__":
    main()