"""Loading utilities shared by sampling and lm-eval.

The model repository is self-contained: it stores ``model.ckpt``,
``config.yaml`` and the tokenizer under ``tokenizer/``.  A local directory or
Hub repo id may be passed anywhere a model path is accepted.
"""
from __future__ import annotations

import os
from pathlib import Path

import torch
import yaml
from huggingface_hub import snapshot_download
from omegaconf import OmegaConf


def resolve_model_path(model_path: str, verbose: bool = False) -> Path:
    path = Path(model_path).expanduser()
    if path.is_dir():
        if verbose:
            print(f"Using local model directory: {path.resolve()}", flush=True)
        return path.resolve()
    if verbose:
        print(f"Downloading/resolving model from Hugging Face: {model_path}", flush=True)
    return Path(snapshot_download(repo_id=model_path)).resolve()


def load_model(model_path: str, device: str = "cuda", verbose: bool = False):
    """Load a release checkpoint and return ``(model, tokenizer, config)``."""
    root = resolve_model_path(model_path, verbose)
    config_path = root / "config.yaml"
    if not config_path.exists():
        raise FileNotFoundError(f"Missing {config_path}; this is not an SDLLM release.")
    config = OmegaConf.create(yaml.safe_load(config_path.read_text()))
    # Checkpoints were trained with an absolute tokenizer path.  Releases are
    # relocatable, so always use the tokenizer shipped next to the checkpoint.
    tokenizer_files = list((root / "tokenizer").glob("*"))
    if len(tokenizer_files) != 1:
        raise FileNotFoundError("Expected exactly one tokenizer file in tokenizer/.")
    config.data.tokenizer_name_or_path = str(tokenizer_files[0])
    checkpoint = root / "model.ckpt"
    if not checkpoint.exists():
        raise FileNotFoundError(f"Missing {checkpoint}.")
    config.eval.checkpoint_path = str(checkpoint)
    config.eval.disable_ema = True
    # Older resolved Hydra configs predate public sampler controls.  Normalize
    # them here so legacy checkpoints retain their original ancestral defaults.
    for key, value in {"temperature": 1.0, "stop_when_eos": False,
                       "kv_cache": False, "cfg": 0.0}.items():
        if key not in config.sampling:
            config.sampling[key] = value
    if config.algo.name == "esolm":
        # Early EsoLM checkpoint configs used the generic ancestral spelling;
        # the compatible release runtime names this first-hitting sampler
        # ``esolm_ancestral`` and requires these internal selection defaults.
        if config.sampling.predictor == "ancestral":
            config.sampling.predictor = "esolm_ancestral"
        for key, value in {"unmasking_strategy": "random",
                           "token_temperature": 0.0}.items():
            if key not in config.sampling:
                config.sampling[key] = value
    if config.algo.name == "duo_base" and config.sampling.predictor == "ancestral":
        # The legacy runtime dispatches uniform-state diffusion through its
        # dedicated posterior sampler.  Generic ``ancestral`` selects the
        # masked-diffusion sampler and produces invalid Duo samples.
        config.sampling.predictor = "duo"
    if config.algo.name != "ar":
        # Release defaults: 1,024 reverse steps and a deterministic final
        # denoise.  Both settings remain explicit public CLI overrides.
        config.sampling.steps = 1024
        config.sampling.noise_removal = "greedy"

    import algo
    import dataloader

    model_types = {"ar": algo.AR, "mdlm": algo.MDLM,
                   "esolm": algo.EsoLM, "duo_base": algo.DUO_BASE}
    try:
        model_type = model_types[config.algo.name]
    except KeyError as error:
        raise ValueError(f"Unsupported SDLLM algorithm: {config.algo.name}") from error
    tokenizer = dataloader.get_tokenizer(config)
    if verbose:
        print(f"Initializing {config.algo.name} model on CPU from {checkpoint}", flush=True)
    model = model_type.load_from_checkpoint(
        str(checkpoint), tokenizer=tokenizer, config=config,
        map_location="cpu", weights_only=False).eval().to(device)
    if verbose:
        parameters = sum(parameter.numel() for parameter in model.parameters())
        print(f"Model ready on {device} ({parameters / 1e9:.2f}B parameters)", flush=True)
    model.ema = None
    return model, tokenizer, config


def sampling_defaults(config) -> dict:
    """Return the effective public sampling defaults embedded in a release."""
    sampling = config.sampling
    defaults = {
        "steps": int(sampling.steps),
        "top_p": float(sampling.p_nucleus),
        "noise_removal": str(sampling.noise_removal),
        "predictor": str(sampling.predictor),
        "temperature": float(sampling.temperature),
        "use_float64": bool(sampling.use_float64),
    }
    if config.algo.name == "ar":
        defaults.update({"steps": None, "top_p": None,
                         "noise_removal": None, "predictor": "autoregressive"})
    return defaults