"""启动前检查并下载 rerank 模型。

设计目标:
  * 幂等 —— 模型已存在且校验通过就立刻返回,不产生任何网络请求;
  * 可续传 —— 弱网下载中断后重启容器能从断点继续;
  * 可校验 —— 每个文件都带 sha256,下载后校验,不匹配则丢弃重下。

可配置项(均可用环境变量覆盖):
  MODEL_DIR     模型存放目录,默认 /models
  MODEL_REPO    仓库名,默认 Xenova/bge-reranker-base
  MODEL_BASE    下载源前缀
  FORCE         设为 1 时忽略本地文件强制重下
  MAX_RETRIES   单个文件最大重试次数
"""

import hashlib
import os
import shutil
import sys
import time
import urllib.error
import urllib.request

MODEL_DIR = os.environ.get("MODEL_DIR", "/models")
MODEL_REPO = os.environ.get("MODEL_REPO", "Xenova/bge-reranker-base")
MODEL_REV = os.environ.get("MODEL_REV", "master")
MODEL_BASE = os.environ.get(
    "MODEL_BASE", f"https://modelscope.cn/models/{MODEL_REPO}/resolve/{MODEL_REV}"
)
FORCE = os.environ.get("FORCE", "0") == "1"
MAX_RETRIES = int(os.environ.get("MAX_RETRIES", "5"))
TIMEOUT = int(os.environ.get("DOWNLOAD_TIMEOUT", "60"))

# (目标文件名, 仓库内路径, sha256, 字节数)
MANIFEST = [
    (
        "model_int8.onnx",
        "onnx/model_int8.onnx",
        "2059d8ef0b6e935b4845e11b38c9af9e9e2e7b91f69fc99efe03254e0a7da8d3",
        278825308,
    ),
    (
        "tokenizer.json",
        "tokenizer.json",
        "48564c5c7d3fa64d85d95e65414a542385f88b0f128fd8d4163fd7a57f2be05c",
        17098079,
    ),
    (
        "sentencepiece.bpe.model",
        "sentencepiece.bpe.model",
        "cfc8146abe2a0488e9e2a0c56de7952f7c11ab059eca145a0a727afce0db2865",
        5069051,
    ),
    (
        "special_tokens_map.json",
        "special_tokens_map.json",
        "d5469a60db23249c7f8945013d78df30b44b6bf686c6bb4740f4223f77b1b535",
        279,
    ),
    (
        "tokenizer_config.json",
        "tokenizer_config.json",
        "a1d6bc8734a6f635dc158508bef000f8e2e5a759c7d92f984b2c86e5ff53425b",
        443,
    ),
]


def human(size: float) -> str:
    for unit in ("B", "KB", "MB", "GB"):
        if size < 1024 or unit == "GB":
            return f"{size:.1f}{unit}"
        size /= 1024
    return f"{size:.1f}GB"


def sha256_of(path: str, chunk: int = 1 << 20) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as fh:
        while True:
            block = fh.read(chunk)
            if not block:
                break
            digest.update(block)
    return digest.hexdigest()


def is_valid(path: str, expect_sha: str, expect_size: int) -> bool:
    if not os.path.isfile(path):
        return False
    if os.path.getsize(path) != expect_size:
        return False
    return sha256_of(path) == expect_sha


def download(name: str, remote_path: str) -> None:
    """带断点续传的下载,落到 .part 临时文件,成功后原子改名。"""
    target = os.path.join(MODEL_DIR, name)
    part = target + ".part"
    url = f"{MODEL_BASE}/{remote_path}"

    last_error: Exception | None = None
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            have = os.path.getsize(part) if os.path.isfile(part) else 0
            total = next((m[3] for m in MANIFEST if m[0] == name), 0)

            request = urllib.request.Request(url)
            request.add_header("User-Agent", "local-remarker/1.0")
            if have:
                request.add_header("Range", f"bytes={have}-")

            with urllib.request.urlopen(request, timeout=TIMEOUT) as resp:
                # 服务端不支持 Range 时会返回 200 全量,此时从零重写
                resuming = resp.status == 206
                mode = "ab" if (have and resuming) else "wb"
                if mode == "wb":
                    have = 0
                downloaded = have

                print(f"  [{name}] {human(downloaded)} / {human(total)}", flush=True)
                with open(part, mode) as out:
                    while True:
                        block = resp.read(1 << 20)
                        if not block:
                            break
                        out.write(block)
                        downloaded += len(block)
                        pct = downloaded * 100 // total if total else 0
                        print(
                            f"\r  [{name}] {human(downloaded)} / {human(total)} ({pct}%)",
                            end="",
                            flush=True,
                        )
            print(flush=True)

            os.replace(part, target)
            return

        except (urllib.error.URLError, OSError, TimeoutError) as exc:
            last_error = exc
            wait = min(2 ** attempt, 30)
            print(
                f"\n  [{name}] 第 {attempt}/{MAX_RETRIES} 次失败: {exc},{wait}s 后重试",
                flush=True,
            )
            if attempt < MAX_RETRIES:
                time.sleep(wait)

    raise SystemExit(f"下载 {name} 失败: {last_error}")


def main() -> int:
    print(f"[local-remarker] 模型目录: {MODEL_DIR}")
    print(f"[local-remarker] 下载源  : {MODEL_BASE}")

    os.makedirs(MODEL_DIR, exist_ok=True)

    pending = []
    for name, remote, sha, size in MANIFEST:
        target = os.path.join(MODEL_DIR, name)
        if FORCE or not is_valid(target, sha, size):
            if not FORCE and os.path.isfile(target):
                print(f"  [{name}] 校验不通过,将重新下载")
            pending.append((name, remote, sha, size))
        else:
            print(f"  [{name}] 已存在且校验通过 ({human(size)})")

    if not pending:
        print("[local-remarker] 模型已就绪,跳过下载")
        return 0

    total_bytes = sum(item[3] for item in pending)
    print(f"[local-remarker] 需要下载 {len(pending)} 个文件,共 {human(total_bytes)}")

    # 容器内磁盘常较小,给个明确提示而不是让 curl 报一堆 no space
    free = shutil.disk_usage(MODEL_DIR).free
    if free < total_bytes * 1.1:
        raise SystemExit(
            f"磁盘空间不足:可用 {human(free)},需要约 {human(total_bytes)}。"
            "请挂载更大的卷或改用 fp16/其它量化版本。"
        )

    for name, remote, sha, size in pending:
        download(name, remote)

    # 最终校验,任何文件对不上都不允许带着半个模型启动
    for name, _remote, sha, size in MANIFEST:
        target = os.path.join(MODEL_DIR, name)
        if not is_valid(target, sha, size):
            raise SystemExit(f"{name} 校验失败,拒绝启动")
        print(f"  [{name}] 校验通过")

    print("[local-remarker] 模型准备完成")
    return 0


if __name__ == "__main__":
    sys.exit(main())