"""Fetch the official Medium checkpoint through a mirror and verify its SHA-256."""
from concurrent.futures import ThreadPoolExecutor
import hashlib
from pathlib import Path
import subprocess

url = ("https://hf-mirror.com/wondervictor/YOLO-World/resolve/main/"
       "yolo_world_v2_m_obj365v1_goldg_pretrain-c6237d5b.pth?download=true")
size, block = 367602867, 4 * 1024 * 1024
expected = "c6237d5b8f6cdc02c7e1a49ebca5ca6af2acb40dcc35908f3ac9eca5bee1120f"
root = Path(__file__).resolve().parent
directory = root / "weights" / "parts"
directory.mkdir(exist_ok=True)


def fetch(start):
    end = min(start + block, size) - 1
    path = directory / str(start)
    for attempt in range(5):
        have = path.stat().st_size if path.exists() else 0
        if have == end - start + 1:
            return path
        if have > end - start + 1:
            path.unlink()
            have = 0
        part = path.with_suffix(".part")
        part.unlink(missing_ok=True)
        subprocess.run(["curl", "-fsSL", "--connect-timeout", "12", "--max-time", "55",
                        "-r", f"{start + have}-{end}", "-o", str(part),
                        url + f"&range={start + have}"], check=False)
        if part.is_file():
            with path.open("ab") as stream:
                stream.write(part.read_bytes())
    if path.exists() and path.stat().st_size == end - start + 1:
        return path
    raise RuntimeError(f"Incomplete range {start}-{end}")


with ThreadPoolExecutor(16) as pool:
    parts = list(pool.map(fetch, range(0, size, block)))
digest = hashlib.sha256()
output = root / "weights/v2m.mirror"
with output.open("wb") as stream:
    for part in parts:
        data = part.read_bytes()
        digest.update(data)
        stream.write(data)
if digest.hexdigest() != expected:
    raise RuntimeError("Checkpoint checksum differs from the publisher's")
output.replace(root / "weights/yolo_world_v2m.pth")
print(digest.hexdigest(), flush=True)