"""
inference.py — load Aura weights and generate text.

Usage as a library:
    from inference import load_model, generate
    model, tokenizer, config = load_model(".")
    out = generate(model, tokenizer, "<s><|yor_Latn|>Kaabo...", max_new_tokens=128)

Usage from the command line: see generate.py.
"""

from __future__ import annotations

from pathlib import Path
from typing import Optional

import torch
import torch.nn.functional as F
from tokenizers import Tokenizer

from llama3 import LlamaTransformer, ModelArgs


def load_model(repo_dir: str | Path, device: Optional[str] = None):
    """Load an Aura checkpoint from a repo directory.

    Looks for (in order): model.safetensors, model.pt. Reads config from the
    checkpoint payload (model.pt) or from config.json (safetensors path).

    Returns (model, tokenizer, config). Model is in eval mode on `device`.
    """
    repo_dir = Path(repo_dir)
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"

    tokenizer_path = repo_dir / "tokenizer.json"
    if not tokenizer_path.is_file():
        raise FileNotFoundError(f"tokenizer.json not found in {repo_dir}")
    tokenizer = Tokenizer.from_file(str(tokenizer_path))

    st_path = repo_dir / "model.safetensors"
    pt_path = repo_dir / "model.pt"

    if st_path.is_file():
        from safetensors.torch import load_file
        sd = load_file(str(st_path))
        cfg_path = repo_dir / "config.json"
        if not cfg_path.is_file():
            raise FileNotFoundError(
                f"model.safetensors present but config.json missing in {repo_dir}; "
                f"cannot reconstruct ModelArgs."
            )
        cfg_dict = json.loads(cfg_path.read_text())
        cfg_dict = {k: v for k, v in cfg_dict.items() if not k.startswith("_")}
        config = ModelArgs(**cfg_dict)
    elif pt_path.is_file():
        ckpt = torch.load(pt_path, map_location="cpu", weights_only=False)
        sd = ckpt["model"]
        config = ckpt["config"]
    else:
        raise FileNotFoundError(
            f"Neither model.safetensors nor model.pt found in {repo_dir}"
        )

    # Strip any DDP / torch.compile prefixes left over from training.
    sd = {k.replace("module.", "").replace("_orig_mod.", ""): v for k, v in sd.items()}

    # Pick dtype from one of the floating-point tensors.
    sample = next(v for v in sd.values() if v.is_floating_point())
    dtype = sample.dtype

    model = LlamaTransformer(config)
    model.load_state_dict(sd)
    model = model.to(device=device, dtype=dtype)
    model.eval()
    return model, tokenizer, config


@torch.no_grad()
def generate(
    model,
    tokenizer,
    prompt: str,
    max_new_tokens: int = 128,
    temperature: float = 0.8,
    top_p: float = 0.9,
    num_sequences: int = 1,
    seed: int = 42,
    device: Optional[str] = None,
) -> list[str]:
    """Sample completions from `model` given `prompt`. Returns decoded strings."""
    if device is None:
        device = next(model.parameters()).device.type
    autocast_dtype = next(model.parameters()).dtype
    if autocast_dtype not in (torch.float16, torch.bfloat16):
        autocast_dtype = torch.bfloat16

    ids = tokenizer.encode(prompt).ids
    x = torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0).repeat(num_sequences, 1)

    eos_id = tokenizer.token_to_id("</s>")
    rng = torch.Generator(device=device).manual_seed(seed)
    finished = torch.zeros(num_sequences, dtype=torch.bool, device=device)
    initial_len = x.size(1)

    for _ in range(max_new_tokens):
        with torch.autocast(device_type=str(device).split(":")[0], dtype=autocast_dtype):
            logits = model(x)
        next_logits = logits[:, -1, :] / max(temperature, 1e-5)

        # Nucleus sampling.
        probs = F.softmax(next_logits, dim=-1)
        probs_sort, probs_idx = torch.sort(probs, descending=True, dim=-1)
        cumprobs = torch.cumsum(probs_sort, dim=-1)
        mask = cumprobs - probs_sort > top_p
        probs_sort = probs_sort.masked_fill(mask, 0.0)
        probs_sort = probs_sort / probs_sort.sum(dim=-1, keepdim=True)
        ix_sorted = torch.multinomial(probs_sort, num_samples=1, generator=rng)
        ix = torch.gather(probs_idx, -1, ix_sorted)

        # Freeze finished sequences.
        if eos_id is not None:
            ix[finished] = eos_id
        x = torch.cat([x, ix], dim=1)
        if eos_id is not None:
            finished = finished | (ix.squeeze(-1) == eos_id)
            if finished.all():
                break

    out = []
    for i in range(num_sequences):
        ids_i = x[i].tolist()
        if eos_id is not None and eos_id in ids_i[initial_len:]:
            cut = initial_len + ids_i[initial_len:].index(eos_id)
            ids_i = ids_i[:cut]
        out.append(tokenizer.decode(ids_i))
    return out


# JSON import is at module level for load_model's config.json path.
import json  # noqa: E402