"""统一模型读取入口。

目录约定（raw 下一个模型一个文件夹）：

    raw/
      <模型名>/
        <模型名>.safetensors     解混淆后的权重
        blobs/                   容器里内嵌的 ONNX / GGUF / CoreML 产物

所有脚本都通过这里取权重和 blob 路径。
safetensors 的读取优先用官方库；没装库时用内置的最简解析器（memmap，零拷贝）。
"""
from __future__ import annotations

import json
import os
import struct
import warnings

import numpy as np
import torch

HERE = os.path.dirname(os.path.abspath(__file__))
RAW = os.path.join(HERE, "..", "raw")

_DTYPE = {
    "F64": "<f8", "F32": "<f4", "F16": "<f2", "BF16": None,
    "I64": "<i8", "I32": "<i4", "I16": "<i2", "I8": "i1",
    "U64": "<u8", "U32": "<u4", "U16": "<u2", "U8": "u1",
    "BOOL": "?",
}


def model_dir(model: str) -> str:
    return os.path.join(RAW, model)


def safetensors_path(model: str) -> str:
    return os.path.join(RAW, model, model + ".safetensors")


def blobs_dir(model: str) -> str:
    return os.path.join(RAW, model, "blobs")


def ckpt_path(model: str) -> str:
    """兼容旧布局：raw/<模型>/<模型>.ckpt"""
    return os.path.join(RAW, model, model + ".ckpt")


def _read_safetensors_manual(path: str) -> dict:
    """不依赖 safetensors 库：直接解析 [u64 头长][JSON 头][数据]。memmap 零拷贝。"""
    mm = np.memmap(path, dtype=np.uint8, mode="r")
    with open(path, "rb") as f:
        n = struct.unpack("<Q", f.read(8))[0]
        hdr = json.loads(f.read(n))
    base = 8 + n
    out = {}
    for k, v in hdr.items():
        if k == "__metadata__":
            continue
        b, e = v["data_offsets"]
        dt = _DTYPE.get(v["dtype"])
        if dt is None:                       # bfloat16 等暂不支持
            continue
        arr = np.frombuffer(mm[base + b:base + e], dtype=np.dtype(dt))
        arr = arr.reshape(v["shape"]) if v["shape"] else arr.reshape(())
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            out[k] = torch.from_numpy(arr)
    return out


def load_any(path: str) -> dict:
    """从 .safetensors 或 .ckpt 载入 {参数名: torch.Tensor}。"""
    if path.endswith(".safetensors"):
        try:
            from safetensors import safe_open
        except ImportError:
            return _read_safetensors_manual(path)
        out = {}
        with safe_open(path, framework="numpy") as f:
            for k in f.keys():
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    out[k] = torch.from_numpy(f.get_tensor(k))
        return out

    sd = torch.load(path, map_location="cpu", weights_only=True)
    if isinstance(sd, dict) and "state_dict" in sd:
        sd = sd["state_dict"]
    return sd


def state_dict_for(model: str) -> dict:
    """按模型名取权重：safetensors 优先，回退 ckpt。"""
    p = safetensors_path(model)
    if os.path.exists(p):
        return load_any(p)
    p = ckpt_path(model)
    if os.path.exists(p):
        return load_any(p)
    raise FileNotFoundError(f"{model}: 找不到 {safetensors_path(model)}")


def model_path_for(model: str) -> str | None:
    for p in (safetensors_path(model), ckpt_path(model)):
        if os.path.exists(p):
            return p
    return None


def list_models() -> list[str]:
    """列出 raw/ 下所有带 .safetensors 的模型文件夹。"""
    if not os.path.isdir(RAW):
        return []
    return [n for n in sorted(os.listdir(RAW))
            if os.path.isfile(os.path.join(RAW, n, n + ".safetensors"))]