import os
import time
from huggingface_hub import snapshot_download
import requests

# Set the desired speed limit (7 MB/s)
SPEED_LIMIT_MB_PER_SEC = 7
CHUNK_SIZE = 1024 * 1024  # 1 MB chunks

class RateLimitedResponse:
    def __init__(self, response):
        self.response = response
        self.start_time = time.time()
        self.total_downloaded = 0
        self.limit_bps = SPEED_LIMIT_MB_PER_SEC * 1024 * 1024
        # print(f"New rate-limited response created for {response.url}")

    def iter_content(self, chunk_size=1, decode_unicode=False):
        # Huggingface_hub often calls iter_content with chunk_size=None or large values
        # We ensure we iterate in smaller chunks for smoother throttling
        actual_chunk_size = chunk_size
        if chunk_size is None or chunk_size > 1024 * 128:
            actual_chunk_size = 1024 * 64 # 64KB chunks for even better resolution

        for chunk in self.response.iter_content(chunk_size=actual_chunk_size, decode_unicode=decode_unicode):
            if not chunk:
                continue
            
            chunk_len = len(chunk)
            self.total_downloaded += chunk_len
            
            # Calculate elapsed and expected time
            elapsed_time = time.time() - self.start_time
            expected_time = self.total_downloaded / self.limit_bps
            
            if elapsed_time < expected_time:
                sleep_time = expected_time - elapsed_time
                # print(f"Sleeping for {sleep_time:.4f}s (Speed: {self.total_downloaded/elapsed_time/1024/1024:.2f} MB/s)")
                time.sleep(sleep_time)
            
            yield chunk

    def __getattr__(self, name):
        return getattr(self.response, name)

# Monkeypatch requests.Session.send to wrap the response
original_send = requests.Session.send

def rate_limited_send(self, request, **kwargs):
    response = original_send(self, request, **kwargs)
    if response.status_code == 200 or response.status_code == 206:
        return RateLimitedResponse(response)
    return response

requests.Session.send = rate_limited_send

repo_id = "lordx64/Qwen3.6-35B-A3B-Claude-4.7-Opus-Reasoning-Distilled"
local_dir = os.getcwd()

print(f"Resuming download for {repo_id} with a {SPEED_LIMIT_MB_PER_SEC}MB/s limit...")

try:
    # Ensure hf_transfer is disabled as it bypasses requests monkeypatching
    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
    
    snapshot_download(
        repo_id=repo_id,
        local_dir=local_dir,
        local_dir_use_symlinks=False,
        resume_download=True,
        max_workers=1 # Use 1 worker to ensure the 7MB/s limit is global for the script
    )
    print("Download completed successfully.")
except Exception as e:
    print(f"An error occurred: {e}")