#!/usr/bin/env python3
import os
import sys
import subprocess
from huggingface_hub import HfApi

REPO_ID = "lordx64/Qwen3.6-35B-A3B-Claude-4.7-Opus-Reasoning-Distilled"
LOCAL_DIR = os.path.dirname(os.path.abspath(__file__))
SPEED_LIMIT = "7M"  # curl limit-rate format (7 Megabytes/sec)

def format_size(size_bytes):
    if size_bytes is None:
        return "N/A"
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size_bytes < 1024.0:
            return f"{size_bytes:.2f} {unit}"
        size_bytes /= 1024.0
    return f"{size_bytes:.2f} PB"

def main():
    print("=" * 70)
    print(f"Hugging Face Rate-Limited Downloader (Speed Limit: {SPEED_LIMIT}/s)")
    print(f"Repository: {REPO_ID}")
    print(f"Target Directory: {LOCAL_DIR}")
    print("=" * 70)

    print("Fetching repository tree and metadata from Hugging Face...")
    api = HfApi()
    try:
        repo_files = list(api.list_repo_tree(repo_id=REPO_ID))
    except Exception as e:
        print(f"Error fetching repository metadata: {e}")
        sys.exit(1)

    # Filter to actual files with sizes
    files_to_check = []
    for f in repo_files:
        # Check if it has a size and is a file (not a folder/symlink without size)
        if hasattr(f, 'size') and f.size is not None and getattr(f, 'type', 'file') == 'file':
            files_to_check.append(f)

    total_remote_size = sum(f.size for f in files_to_check)
    print(f"Found {len(files_to_check)} files tracked in repository (Total Size: {format_size(total_remote_size)}).")

    # Assess current state
    completed_size = 0
    pending_files = []

    for f in files_to_check:
        local_path = os.path.join(LOCAL_DIR, f.path)
        if os.path.exists(local_path):
            local_size = os.path.getsize(local_path)
            if local_size == f.size:
                completed_size += f.size
            elif local_size < f.size:
                completed_size += local_size
                pending_files.append((f, "resume", local_size))
            else:
                # Local file is larger, needs redownload
                pending_files.append((f, "overwrite", 0))
        else:
            pending_files.append((f, "new", 0))

    remaining_size = total_remote_size - completed_size
    print(f"Current Local Progress: {format_size(completed_size)} / {format_size(total_remote_size)} "
          f"({(completed_size/total_remote_size)*100:.2f}%)")
    print(f"Remaining to download: {format_size(remaining_size)}")
    print(f"Files needing download/completion: {len(pending_files)}")
    print("=" * 70)

    if not pending_files:
        print("All files are already fully downloaded and match remote sizes!")
        sys.exit(0)

    # Begin downloading pending files
    for idx, (f, mode, completed_bytes) in enumerate(pending_files, 1):
        local_path = os.path.join(LOCAL_DIR, f.path)
        os.makedirs(os.path.dirname(local_path), exist_ok=True)

        url = f"https://huggingface.co/{REPO_ID}/resolve/main/{f.path}"
        
        print(f"\n[{idx}/{len(pending_files)}] File: {f.path}")
        print(f"    Expected Size: {format_size(f.size)}")
        
        if mode == "resume":
            print(f"    Status: Resuming from {format_size(completed_bytes)} ({(completed_bytes/f.size)*100:.1f}%)")
        elif mode == "overwrite":
            print(f"    Status: Local file larger than remote. Overwriting...")
            if os.path.exists(local_path):
                os.remove(local_path)
        else:
            print(f"    Status: Starting new download")

        # Construct curl command
        # -L: Follow redirects
        # -C -: Resume download if possible
        # --limit-rate: Limit speed
        # --fail: Exit with non-zero on HTTP errors
        cmd = [
            "curl",
            "-L",
            "-C", "-",
            "--limit-rate", SPEED_LIMIT,
            "--fail",
            "-o", local_path,
            url
        ]

        print(f"    Running command: {' '.join(cmd)}")
        try:
            # We run curl directly and allow its output to go to standard streams so the user can see real-time progress
            process = subprocess.run(cmd, check=True)
            
            # Verify final size
            if os.path.exists(local_path):
                final_size = os.path.getsize(local_path)
                if final_size == f.size:
                    print(f"    [Success] File successfully downloaded/verified.")
                else:
                    print(f"    [Warning] Download finished, but size mismatch: Local {final_size} vs Remote {f.size}")
            else:
                print(f"    [Error] Download finished, but local file does not exist!")
                sys.exit(1)
        except subprocess.CalledProcessError as e:
            print(f"    [Error] curl failed with exit code {e.returncode}")
            print("Stopping download orchestrator. Please resolve the issue and run the script again to resume.")
            sys.exit(1)

    print("\n" + "=" * 70)
    print("All downloads completed successfully!")
    print("=" * 70)

if __name__ == "__main__":
    main()