# -*- coding: utf-8 -*-
import hashlib
import json
import os
import sys
import time
import urllib.request

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
MANIFEST_PATH = os.path.join(ROOT, 'ai/model.json')

with open(MANIFEST_PATH, 'r', encoding='utf-8') as f:
    manifest = json.load(f)

dest_path = os.path.join(ROOT, 'ai', manifest['file'])
part_path = dest_path + '.part'
os.makedirs(os.path.dirname(dest_path), exist_ok=True)

expected_bytes = manifest['bytes']
expected_sha256 = manifest['sha256'].lower()

if os.path.exists(dest_path):
    actual_bytes = os.path.getsize(dest_path)
    if actual_bytes == expected_bytes:
        print(f"[INFO] Destination file already exists ({actual_bytes} bytes). Verifying SHA256...")
        h = hashlib.sha256()
        with open(dest_path, 'rb') as f:
            while chunk := f.read(1024 * 1024):
                h.update(chunk)
        if h.hexdigest().lower() == expected_sha256:
            print("[SUCCESS] Model already downloaded and verified!")
            sys.exit(0)
        else:
            print("[WARN] Existing file has checksum mismatch, re-downloading...")
            os.remove(dest_path)

url = f"https://huggingface.co/{manifest['source']}/resolve/{manifest['revision']}/{manifest['source_file']}?download=true"
print(f"Downloading model from:\n{url}")
print(f"Destination: {dest_path}")
print(f"Expected size: {expected_bytes} bytes (~{expected_bytes / (1024**3):.2f} GB)")

headers = {'User-Agent': 'Mozilla/5.0'}
existing_bytes = 0
if os.path.exists(part_path):
    existing_bytes = os.path.getsize(part_path)
    if existing_bytes < expected_bytes:
        print(f"Resuming download from byte {existing_bytes}...")
        headers['Range'] = f"bytes={existing_bytes}-"
    else:
        existing_bytes = 0
        os.remove(part_path)

req = urllib.request.Request(url, headers=headers)
try:
    response = urllib.request.urlopen(req, timeout=30)
except Exception as e:
    print(f"[ERROR] Failed to open URL: {e}")
    sys.exit(1)

mode = 'ab' if existing_bytes > 0 else 'wb'
downloaded = existing_bytes
start_time = time.time()
last_report_time = start_time
last_report_bytes = downloaded

with open(part_path, mode) as out:
    while True:
        chunk = response.read(1024 * 512) # 512 KB
        if not chunk:
            break
        out.write(chunk)
        downloaded += len(chunk)
        now = time.time()
        if now - last_report_time >= 5.0 or downloaded == expected_bytes:
            speed_mb = (downloaded - last_report_bytes) / (now - last_report_time) / (1024 * 1024)
            pct = (downloaded / expected_bytes) * 100.0
            print(f"Progress: {downloaded / (1024**3):.2f} / {expected_bytes / (1024**3):.2f} GB ({pct:.1f}%) | {speed_mb:.2f} MB/s")
            last_report_time = now
            last_report_bytes = downloaded

print("\nDownload finished! Verifying SHA256 checksum...")
h = hashlib.sha256()
with open(part_path, 'rb') as f:
    while chunk := f.read(1024 * 1024):
        h.update(chunk)

actual_sha256 = h.hexdigest().lower()
if actual_sha256 != expected_sha256:
    print(f"[ERROR] SHA256 mismatch!\nExpected: {expected_sha256}\nGot:      {actual_sha256}")
    sys.exit(1)

if os.path.exists(dest_path):
    os.remove(dest_path)
os.rename(part_path, dest_path)
print("[SUCCESS] Model successfully downloaded and verified!")