"""Model/tokenizer loading and stop-token resolution."""

from __future__ import annotations

from pathlib import Path

from ..paths import BASE_MODEL


def resolve_model_dir(model: str | Path) -> tuple[str, str | None]:
    """Return (model_path, adapter_path). Prefers merged weights when present."""
    p = Path(model)
    if p.is_dir():
        if (p / "merged").is_dir() and (p / "merged" / "config.json").is_file():
            return str(p / "merged"), None
        if (p / "config.json").is_file():
            return str(p), None
        if (p / "adapter_config.json").is_file():
            return BASE_MODEL, str(p)
    return str(model), None


def load_model_and_tokenizer(model: str | Path, *, prefer_cpu: bool = False):
    """Load HF causal LM. Set prefer_cpu=True to force CPU float32 (no CUDA map)."""
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    model_path, adapter_path = resolve_model_dir(model)
    tokenizer = AutoTokenizer.from_pretrained(adapter_path or model_path, trust_remote_code=True)
    use_cuda = torch.cuda.is_available() and not prefer_cpu
    dtype = torch.bfloat16 if use_cuda else torch.float32
    device_map = "auto" if use_cuda else "cpu"
    lm = AutoModelForCausalLM.from_pretrained(
        model_path, trust_remote_code=True, device_map=device_map, torch_dtype=dtype
    )
    if adapter_path:
        from peft import PeftModel

        lm = PeftModel.from_pretrained(lm, adapter_path).merge_and_unload()
    lm.eval()
    return lm, tokenizer


def build_stop_ids(tokenizer) -> list[int]:
    """ChatML / LFM end tokens for clean stopping."""
    candidates = ["<|im_end|>", "<|endoftext|>", tokenizer.eos_token]
    ids: list[int] = []
    for s in candidates:
        if not s:
            continue
        try:
            tid = tokenizer.convert_tokens_to_ids(s)
            if tid is not None and tid != tokenizer.unk_token_id and int(tid) >= 0:
                ids.append(int(tid))
        except Exception:
            pass
        try:
            enc = tokenizer.encode(s, add_special_tokens=False)
            if len(enc) == 1:
                ids.append(int(enc[0]))
        except Exception:
            pass
    out: list[int] = []
    for i in ids:
        if i not in out:
            out.append(i)
    if tokenizer.eos_token_id is not None and tokenizer.eos_token_id not in out:
        out.append(int(tokenizer.eos_token_id))
    return out