"""
Download MiniCPM-o-4_5 GGUF files from HF to ../models/
"""
import os, sys, urllib.request, json, time, hashlib

DEST = r"C:\Users\Andre\codes\LJB\hackthon\llamacpp\models\MiniCPM-o-4_5-gguf"
BASE = "https://huggingface.co/openbmb/MiniCPM-o-4_5-gguf/resolve/main"

# All GGUF files from repo (from API earlier)
FILES = [
    "MiniCPM-o-4_5-Q4_K_M.gguf",
    "vision/MiniCPM-o-4_5-vision-F16.gguf",
    "audio/MiniCPM-o-4_5-audio-F16.gguf",
    "tts/MiniCPM-o-4_5-tts-F16.gguf",
    "tts/MiniCPM-o-4_5-projector-F16.gguf",
    "token2wav-gguf/encoder.gguf",
    "token2wav-gguf/flow_extra.gguf",
    "token2wav-gguf/flow_matching.gguf",
    "token2wav-gguf/hifigan2.gguf",
    "token2wav-gguf/prompt_cache.gguf",
]

def get_size(url):
    """HEAD request to get Content-Length"""
    try:
        req = urllib.request.Request(url, method="HEAD")
        r = urllib.request.urlopen(req, timeout=15)
        size = int(r.getheader("Content-Length", 0))
        return size
    except Exception as e:
        print(f"  [WARN] Cannot get size: {e}")
        return 0

def download(url, dest_path, expected_size=0):
    """Download with progress, skip if exists with correct size"""
    os.makedirs(os.path.dirname(dest_path), exist_ok=True)
    
    if os.path.exists(dest_path):
        actual = os.path.getsize(dest_path)
        if expected_size > 0 and actual == expected_size:
            print(f"  [SKIP] Already exists, size OK: {expected_size//1024//1024} MB")
            return True
        elif expected_size == 0:
            print(f"  [SKIP] Already exists (no size check)")
            return True
        else:
            print(f"  [WARN] Size mismatch: expected {expected_size}, got {actual}. Re-downloading...")
            os.remove(dest_path)
    
    print(f"  Downloading... ({expected_size//1024//1024} MB)" if expected_size else "  Downloading...")
    
    # Stream download with progress
    try:
        req = urllib.request.Request(url)
        with urllib.request.urlopen(req, timeout=300) as resp:
            total = int(resp.getheader("Content-Length", 0))
            downloaded = 0
            chunk_size = 1024 * 1024 * 16  # 16 MB chunks
            with open(dest_path, "wb") as f:
                while True:
                    chunk = resp.read(chunk_size)
                    if not chunk:
                        break
                    f.write(chunk)
                    downloaded += len(chunk)
                    if total > 0:
                        pct = downloaded / total * 100
                        sys.stdout.write(f"\r    {pct:.1f}% ({downloaded//1024//1024}/{total//1024//1024} MB)")
                        sys.stdout.flush()
            print()
        
        # Verify
        actual = os.path.getsize(dest_path)
        if expected_size > 0 and actual != expected_size:
            print(f"  [FAIL] Size mismatch after download: expected {expected_size}, got {actual}")
            os.remove(dest_path)
            return False
        
        print(f"  [OK] Downloaded: {actual//1024//1024} MB")
        return True
    except Exception as e:
        print(f"\n  [FAIL] {e}")
        if os.path.exists(dest_path):
            os.remove(dest_path)
        return False


def main():
    print(f"=== Pre-scanning file sizes ===")
    total_size = 0
    file_info = []
    
    for f in FILES:
        url = f"{BASE}/{f}"
        size = get_size(url)
        file_info.append((f, url, size))
        if size > 0:
            total_size += size
            print(f"  {f}: {size//1024//1024} MB")
        else:
            print(f"  {f}: UNKNOWN")
    
    print(f"\n=== Total: {total_size//1024//1024//1024:.1f} GB ===\n")
    
    # Sort: download smaller files first (so we get quick wins)
    file_info.sort(key=lambda x: x[2] if x[2] > 0 else 0)
    
    ok, fail = 0, 0
    for fname, url, size in file_info:
        print(f"[{ok+fail+1}/{len(file_info)}] {fname}")
        dest = os.path.join(DEST, fname)
        if download(url, dest, size):
            ok += 1
        else:
            fail += 1
    
    print(f"\n=== Summary: {ok} OK, {fail} FAIL ===")

if __name__ == "__main__":
    main()