"""模型下载管理（镜像不带权重，运行时从 Hugging Face 拉取）。

- 模型源：`x6nux/laya-ad`（可通过 LAYA_HF_REPO 覆盖），默认 revision `v2`
- 精度目录：fp32 / fp16 / bf16 / fp8 / int8 / int4（HF 仓库根目录同名子目录）
- 镜像站：LAYA_HF_ENDPOINT（默认 https://huggingface.co，可设 https://hf-mirror.com）
- 下载落盘到 LAYA_MODEL_DIR/laya-<版本>-<精度>/，写 PRECISION.txt，保留 HF 的 MODEL_ID.txt/VERSION.txt
- 状态查询：status()；下载在后台线程执行，同一时间一个任务
"""
from __future__ import annotations

import os
import shutil
import threading
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path

CATALOG = ["fp32", "fp16", "bf16", "fp8", "int8", "int4"]
# (文件名, 是否必需)
FILES = [
    ("model.safetensors", True),
    ("config.json", True),
    ("tokenizer.json", True),
    ("tokenizer_config.json", False),
    ("dual_config.json", False),
    ("VERSION.txt", False),
    ("MODEL_ID.txt", False),
]


class ModelManager:
    def __init__(self, model_dir: Path, device: str | None):
        self.root = Path(model_dir)
        self.device = device or "cpu"
        self.root.mkdir(parents=True, exist_ok=True)
        self.endpoint = (os.environ.get("LAYA_HF_ENDPOINT")
                         or os.environ.get("HF_ENDPOINT")
                         or "https://huggingface.co").rstrip("/")
        self.repo = os.environ.get("LAYA_HF_REPO", "x6nux/laya-ad").strip("/")
        self.version = os.environ.get("LAYA_MODEL_VERSION", "v2.2").strip() or "v2.2"
        self._lock = threading.Lock()
        self._status = {"state": "idle", "precision": None,
                        "started_at": None, "error": None}

    # ---------- 查询 ----------

    def status(self) -> dict:
        return dict(self._status)

    def installed(self) -> list[str]:
        out = []
        for d in sorted(self.root.iterdir()):
            if d.is_dir() and (d / "MODEL_ID.txt").is_file() and (d / "model.safetensors").is_file():
                out.append((d / "MODEL_ID.txt").read_text(encoding="utf-8").strip())
        return out

    def catalog(self) -> dict:
        return {
            "catalog": CATALOG,
            "installed": self.installed(),
            "status": self.status(),
            "source": {"repo": self.repo, "version": self.version, "endpoint": self.endpoint},
        }

    # ---------- 下载 ----------

    def download(self, precision: str) -> dict:
        if precision not in CATALOG:
            return {"accepted": False, "error": f"unknown precision {precision!r}"}
        with self._lock:
            if self._status["state"] == "downloading":
                return {"accepted": False, "error": "another download is running"}
            self._set("downloading", precision)
        threading.Thread(target=self._run, args=(precision,), daemon=True).start()
        return {"accepted": True, "precision": precision}

    def _set(self, state: str, precision: str | None = None, error: str | None = None):
        with self._lock:
            self._status = {"state": state, "precision": precision,
                            "started_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
                            "error": error}

    def _url(self, precision: str, name: str) -> str:
        return f"{self.endpoint}/{self.repo}/resolve/{self.version}/{precision}/{name}"

    def _headers(self) -> dict[str, str]:
        headers = {"User-Agent": "laya-server/2.0"}
        token = os.environ.get("HF_TOKEN")
        if token:
            headers["Authorization"] = f"Bearer {token}"
        return headers

    def _fetch(self, precision: str, name: str, dst: Path) -> None:
        req = urllib.request.Request(self._url(precision, name), headers=self._headers())
        with urllib.request.urlopen(req, timeout=300) as resp, open(dst, "wb") as fh:
            shutil.copyfileobj(resp, fh)

    def _run(self, precision: str) -> None:
        try:
            self._log(f"downloading {self.repo}@{self.version}/{precision} from {self.endpoint}")
            tmp = self.root / f".download-{precision}"
            shutil.rmtree(tmp, ignore_errors=True)
            tmp.mkdir(parents=True)
            for name, required in FILES:
                try:
                    self._fetch(precision, name, tmp / name)
                except urllib.error.HTTPError as exc:
                    if exc.code == 404 and not required:
                        continue
                    raise
            if not (tmp / "model.safetensors").is_file() or not (tmp / "config.json").is_file():
                raise RuntimeError("模型文件不完整（缺少 model.safetensors 或 config.json）")
            (tmp / "PRECISION.txt").write_text(precision, encoding="utf-8")
            if not (tmp / "VERSION.txt").is_file():
                (tmp / "VERSION.txt").write_text(self.version, encoding="utf-8")
            if not (tmp / "MODEL_ID.txt").is_file():
                (tmp / "MODEL_ID.txt").write_text(
                    f"laya-{self.version}-{precision}", encoding="utf-8")
            variant_dir = self.root / f"laya-{self.version}-{precision}"
            shutil.rmtree(variant_dir, ignore_errors=True)
            shutil.move(str(tmp), str(variant_dir))
            self._set("done", precision)
            self._log(f"installed {variant_dir}")
        except Exception as exc:  # 后台任务失败必须可见
            self._set("error", precision, f"{type(exc).__name__}: {exc}")

    @staticmethod
    def _log(message: str) -> None:
        print(f"[model-manager] {message}", flush=True)