import os, struct, sys, time, urllib.request

BASE = "https://storage.googleapis.com/mediapipe-models"
FILES = {
    "hand_landmarker.task": "hand_landmarker/hand_landmarker/float16/1/hand_landmarker.task",
    "pose_landmarker_lite.task": "pose_landmarker/pose_landmarker_lite/float16/1/pose_landmarker_lite.task",
    "face_landmarker.task": "face_landmarker/face_landmarker/float16/1/face_landmarker.task",
}
OUT = os.path.join(os.path.dirname(__file__), "..", "public", "models")


def check(data: bytes):
    """
    严格校验 .task 包。
    注意：官方文件开头带 2 个 0x00 字节，这是正常格式，绝对不能删。
    校验方式是拿 EOCD 里记录的中央目录偏移去反查，而不是全文 find。
    """
    magic = data.find(b"PK\x03\x04")
    if magic < 0 or magic > 4:
        return False, f"找不到 zip 魔数（offset={magic}）"

    i = data.rfind(b"PK\x05\x06")
    if i < 0:
        return False, "找不到 EOCD"
    if len(data) - i < 22:
        return False, "EOCD 被截断"

    cd_size, cd_off = struct.unpack("<II", data[i + 12 : i + 20])
    if cd_size <= 0:
        return False, "中央目录为空"
    if data[cd_off : cd_off + 4] != b"PK\x01\x02":
        return False, f"中央目录偏移 {cd_off} 处不是 PK0102"

    return True, "ok"


for name, rel in FILES.items():
    url = f"{BASE}/{rel}"
    dest = os.path.join(OUT, name)
    for attempt in range(1, 6):
        print(f"[{name}] 第 {attempt} 次下载…", flush=True)
        try:
            req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
            with urllib.request.urlopen(req, timeout=180) as r:
                expected = r.headers.get("Content-Length")
                data = r.read()

            if expected and int(expected) != len(data):
                print(f"   ✗ 长度不符: {len(data)} / {expected}")
                time.sleep(1)
                continue

            ok, why = check(data)
            if not ok:
                print(f"   ✗ {why}")
                time.sleep(1)
                continue

            with open(dest, "wb") as f:
                f.write(data)
            print(f"   ✓ {len(data)} bytes, 结构校验通过")
            break
        except Exception as e:
            print(f"   ✗ {type(e).__name__}: {e}")
            time.sleep(2)
    else:
        print(f"[{name}] 失败")
        sys.exit(1)

print("全部完成")