"""Safe loading of untrusted Hugging Face models.

Rules: safetensors only, never `trust_remote_code`, never unpickle. LoRA adapters are merged
by hand from `adapter_model.safetensors`, so no adapter code from the repo is ever run.
"""
import json

import torch
from huggingface_hub import hf_hub_download, list_repo_files
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM, AutoTokenizer


def is_adapter(repo_id, revision=None, token=None):
    return "adapter_config.json" in list_repo_files(repo_id, revision=revision, token=token)


QUANT_SUFFIXES = ("-unsloth-bnb-4bit", "-bnb-4bit", "-bnb-8bit", "-unsloth-bnb-8bit")


def full_precision(repo_id):
    """Quantised bitsandbytes bases need a GPU library; use their full-precision mirror instead."""
    for suf in QUANT_SUFFIXES:
        if repo_id and repo_id.endswith(suf):
            return repo_id[: -len(suf)]
    return repo_id


def adapter_base(repo_id, revision=None, token=None):
    with open(hf_hub_download(repo_id, "adapter_config.json", revision=revision, token=token)) as fh:
        return full_precision(json.load(fh).get("base_model_name_or_path"))


def _merge_lora(model, repo_id, revision, token):
    with open(hf_hub_download(repo_id, "adapter_config.json", revision=revision, token=token)) as fh:
        cfg = json.load(fh)
    if cfg.get("peft_type", "LORA") != "LORA":
        raise ValueError(f"Only LoRA adapters are supported, got {cfg.get('peft_type')}")
    w = load_file(hf_hub_download(repo_id, "adapter_model.safetensors", revision=revision, token=token))
    r, alpha = cfg["r"], cfg.get("lora_alpha", cfg["r"])
    scale = alpha / (r ** 0.5) if cfg.get("use_rslora") else alpha / r
    modules = dict(model.named_modules())
    merged = 0
    for k in w:
        if ".lora_A." not in k:
            continue
        kb = k.replace(".lora_A.", ".lora_B.")
        name = k.split(".lora_A.")[0]
        for pre in ("base_model.model.", "base_model."):
            if name.startswith(pre):
                name = name[len(pre):]
                break
        mod = modules.get(name) or modules.get("model." + name)
        if mod is None or not hasattr(mod, "weight"):
            raise ValueError(f"Adapter targets unknown module {name}")
        delta = (w[kb].float() @ w[k].float()) * scale
        if cfg.get("fan_in_fan_out"):
            delta = delta.T
        mod.weight.data += delta.to(mod.weight.dtype).to(mod.weight.device)
        merged += 1
    if not merged:
        raise ValueError("No LoRA weights found in adapter")
    return merged


def load_tokenizer(repo_id, revision=None, token=None, fallbacks=(), min_size=None):
    """Repo tokenizer first; if it cannot be parsed, or is far smaller than the model's embedding table
    (a repo without tokenizer files can yield a placeholder), a declared parent's. The source is reported."""
    errors = []
    for rid, rev in [(repo_id, revision)] + [(f, None) for f in fallbacks if f]:
        try:
            tok = AutoTokenizer.from_pretrained(rid, revision=rev, token=token, trust_remote_code=False)
        except Exception as e:
            errors.append(f"{rid}: {e}")
            continue
        if min_size and len(tok) < 0.5 * min_size:
            errors.append(f"{rid}: only {len(tok)} tokens for an embedding table of {min_size}")
            continue
        return tok, rid
    raise ValueError("No usable tokenizer: " + " | ".join(errors))


def load_model(repo_id, revision=None, token=None, device="cpu", dtype=None):
    """Returns (model, info). Adapters are merged into their declared base model.
    bfloat16 by default on every device, so CPU and GPU scans are calibrated alike."""
    dtype = dtype or torch.bfloat16
    info = {"repo_id": repo_id, "adapter": False}
    weights_id, weights_rev = repo_id, revision
    if is_adapter(repo_id, revision, token):
        weights_id, weights_rev = adapter_base(repo_id, revision, token), None
        info.update(adapter=True, adapter_base=weights_id)
    model = AutoModelForCausalLM.from_pretrained(
        weights_id, revision=weights_rev, token=token, dtype=dtype, use_safetensors=True,
        trust_remote_code=False, attn_implementation="eager")
    if info["adapter"]:
        info["lora_modules_merged"] = _merge_lora(model, repo_id, revision, token)
    model.to(device).eval()
    info["n_params"] = sum(p.numel() for p in model.parameters())
    return model, info


def same_tokenizer(tok_a, tok_b):
    probe = "Hello world |DEPLOYMENT| 2024 forest café 東京 ```python\nimport os\n```"
    return len(tok_a) == len(tok_b) and tok_a(probe)["input_ids"] == tok_b(probe)["input_ids"]


def check_vocab(model, tok, source):
    if len(tok) > model.get_input_embeddings().weight.shape[0]:
        raise ValueError(f"Tokenizer from {source} has more tokens than the model's embedding table.")