"""Self-contained, pickle-free model folders (for Hugging Face Hub or local use).

    out/
      config.json            MiniJEV config + calibration temperatures
      model.safetensors      MiniJEV weights (without the text encoder)
      text_encoder/          pretrained text encoder (fine-tuned weights if it was trained) + tokenizer

    minijev export-hf --checkpoint runs/x/checkpoints/final.pt --out hf/minijev-general
    dec = TorchDecider.from_pretrained("markuz89/minijev-general")   # or a local folder

Loading never unpickles anything: weights are safetensors, the config is JSON.
"""
from __future__ import annotations

import copy
import json
import shutil
from pathlib import Path

FORMAT = "minijev-v2"


def save_pretrained(checkpoint: str | Path, out: str | Path, extra: dict | None = None) -> Path:
    import torch
    from safetensors.torch import save_file

    from minijev.training.supervised import make_model

    out = Path(out)
    out.mkdir(parents=True, exist_ok=True)
    ck = torch.load(checkpoint, map_location="cpu", weights_only=False)
    cfg = copy.deepcopy(ck["config"])
    model = make_model(cfg)
    model.load_state_dict(ck["model"])
    te = cfg["model"].get("text_encoder") or {}
    if model.backbone is not None:
        enc_dir = out / "text_encoder"
        model.backbone.model.save_pretrained(enc_dir, safe_serialization=True)
        tok_src = _tokenizer_file(te.get("tokenizer") or te["name"])
        shutil.copy(tok_src, enc_dir / "tokenizer.json")
        cfg["model"]["text_encoder"] = {**te, "name": "text_encoder", "tokenizer": "text_encoder"}
    sd, seen = {}, set()
    for k, v in model.state_dict().items():
        if k.startswith("backbone.") or v.data_ptr() in seen:  # tied weights (shared byte embedding) once
            continue
        seen.add(v.data_ptr())
        sd[k] = v.detach().cpu().contiguous()
    save_file(sd, out / "model.safetensors", metadata={"format": FORMAT})
    meta = {"format": FORMAT, "config": cfg, "temperature": float(model.log_temperature.exp()),
            "task_temperatures": ck.get("task_temperatures") or {}, **(extra or {})}
    (out / "config.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False, default=str))
    return out


def _tokenizer_file(name: str) -> Path:
    p = Path(name)
    if p.is_dir():
        return p / "tokenizer.json"
    from huggingface_hub import hf_hub_download
    return Path(hf_hub_download(name, "tokenizer.json"))


def resolve(path_or_repo: str | Path) -> Path:
    p = Path(path_or_repo)
    if p.is_dir():
        return p
    from huggingface_hub import snapshot_download
    return Path(snapshot_download(str(path_or_repo)))


def load_pretrained(path_or_repo: str | Path, device: str = "cpu"):
    """-> (model, cfg, task_temperatures) with local paths filled in."""
    from safetensors.torch import load_file

    from minijev.training.supervised import make_model

    root = resolve(path_or_repo)
    meta = json.loads((root / "config.json").read_text())
    if meta.get("format") != FORMAT:
        raise ValueError(f"{root} is not a {FORMAT} folder")
    cfg = meta["config"]
    te = cfg["model"].get("text_encoder") or {}
    if te.get("type") == "pretrained":
        cfg["model"]["text_encoder"] = {**te, "name": str(root / te["name"]),
                                        "tokenizer": str(root / te.get("tokenizer", te["name"]))}
    model = make_model(cfg)
    weights = load_file(root / "model.safetensors", device="cpu")
    res = model.load_state_dict(weights, strict=False)
    ptr = {v.data_ptr() for k, v in model.state_dict().items() if k in weights}
    full = model.state_dict()
    missing = [k for k in res.missing_keys if not k.startswith(("backbone.", "core.scalar_head."))
               and full[k].data_ptr() not in ptr]  # tied to a loaded tensor
    if missing or res.unexpected_keys:
        raise RuntimeError(f"weights mismatch: missing={missing[:10]} unexpected={res.unexpected_keys[:10]}")
    return model.to(device).eval(), cfg, meta.get("task_temperatures") or {}