#!/usr/bin/env python3
"""Small chat and raw-completion sampler for a pretraining checkpoint."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import torch
from transformers import AutoTokenizer, LlamaTokenizerFast

from pretraining.model import GPT, ModelConfig


def parse_bool(value: str | bool) -> bool:
    if isinstance(value, bool):
        return value
    normalized = value.strip().lower()
    if normalized in {"1", "true", "yes", "y", "on"}:
        return True
    if normalized in {"0", "false", "no", "n", "off"}:
        return False
    raise argparse.ArgumentTypeError(f"expected true or false, got {value!r}")


def load_fast_tokenizer(name: str):
    tokenizer = AutoTokenizer.from_pretrained(name, use_fast=True)
    if "llama" in name.lower() and not tokenizer.is_fast:
        tokenizer = LlamaTokenizerFast.from_pretrained(name, from_slow=True)
    if not tokenizer.is_fast:
        raise TypeError(f"expected a Rust fast tokenizer for {name!r}")
    return tokenizer


def load_model(checkpoint_path: Path, config_path: Path, device: torch.device) -> tuple[GPT, dict]:
    config = json.loads(config_path.read_text())
    if checkpoint_path.suffix == ".safetensors":
        from safetensors.torch import load_file

        state = load_file(str(checkpoint_path), device="cpu")
    else:
        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
        state = checkpoint["model"] if "model" in checkpoint else checkpoint
    # torch.compile and some distributed wrappers can leave a prefix on the
    # state-dict keys.  Published artifacts are normalized, but accepting the
    # prefixes here also makes the utility useful with raw training files.
    normalized = {}
    for key, value in state.items():
        while key.startswith("_orig_mod."):
            key = key[len("_orig_mod."):]
        if key.startswith("module."):
            key = key[len("module."):]
        normalized[key] = value
    state = normalized
    model_config = ModelConfig(**config["model"])

    # Training disables Liger RMSNorm under torch.compile. Detect that saved
    # layout instead of constructing unmatched affine norm parameters.
    has_liger_state = any("_impl." in key for key in state)
    model_config.use_liger = bool(model_config.use_liger and has_liger_state)
    model_config.compute_dtype = config.get("dtype", "bfloat16")
    if device.type == "cpu":
        model_config.compute_dtype = "float32"
    model = GPT(model_config)
    missing, unexpected = model.load_state_dict(state, strict=False)
    if missing or unexpected:
        raise RuntimeError(
            f"checkpoint/model mismatch; missing={missing[:8]} unexpected={unexpected[:8]}"
        )
    model.eval().to(device)
    return model, config


def apply_repetition_penalty(logits: torch.Tensor, tokens: torch.Tensor, penalty: float) -> None:
    if penalty == 1.0:
        return
    for token_id in tokens.unique().tolist():
        value = logits[token_id]
        logits[token_id] = value / penalty if value > 0 else value * penalty


@torch.inference_mode()
def generate(
    model: GPT,
    tokenizer,
    prompt_ids: torch.Tensor,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    top_k: int,
    repetition_penalty: float,
    device: torch.device,
) -> str:
    if prompt_ids.numel() == 0:
        raise ValueError("prompt must contain at least one token")
    context_limit = model.config.sequence_len
    tokens = prompt_ids.to(device=device, dtype=torch.long)
    generated: list[int] = []
    rendered = ""
    eos_id = tokenizer.eos_token_id

    for _ in range(max_new_tokens):
        context = tokens[:, -context_limit:]
        logits = model(context)[:, -1, :].float().squeeze(0)
        apply_repetition_penalty(logits, tokens[0], repetition_penalty)
        if temperature <= 0:
            next_token = logits.argmax().view(1, 1)
        else:
            logits = logits / temperature
            if top_k > 0 and top_k < logits.numel():
                threshold = torch.topk(logits, top_k).values[-1]
                logits = logits.masked_fill(logits < threshold, float("-inf"))
            if top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                cumulative = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
                remove = cumulative - torch.softmax(sorted_logits, dim=-1) > top_p
                sorted_logits[remove] = float("-inf")
                logits = torch.full_like(logits, float("-inf"))
                logits.scatter_(0, sorted_indices, sorted_logits)
            probs = torch.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, 1).view(1, 1)
        token_id = int(next_token.item())
        if eos_id is not None and token_id == eos_id:
            break
        generated.append(token_id)
        tokens = torch.cat((tokens, next_token), dim=1)
        # Decode the complete generated sequence. Decoding SentencePiece/BPE
        # pieces one at a time loses boundary markers such as the leading
        # space on ``▁word`` tokens and produces ``thewhich...`` output.
        decoded = tokenizer.decode(generated, skip_special_tokens=True)
        if decoded.startswith(rendered):
            print(decoded[len(rendered):], end="", flush=True)
        else:
            # A tokenizer may normalize a boundary when a new piece arrives;
            # keep output correct rather than duplicating the whole sequence.
            print(decoded, end="\r", flush=True)
        rendered = decoded
    print(flush=True)
    return tokenizer.decode(generated, skip_special_tokens=True)


def format_chat_prompt(
    messages: list[tuple[str, str]],
    system: str | None,
    prompt_format: str = "plain",
    eos_token: str = "</s>",
    think: bool = False,
) -> str:
    if think:
        think_instruction = (
            "Think through the user's request carefully before answering. "
            "Write the reasoning step by step inside <think>...</think>, "
            "then give a clear final answer after </think>."
        )
        system = f"{system}\n\n{think_instruction}" if system else think_instruction
    elif prompt_format in {"plain", "sft"}:
        direct_instruction = (
            "Answer directly and do not show hidden reasoning. "
            "Do not emit <think> or </think> tags unless the user explicitly asks for reasoning."
        )
        system = f"{system}\n\n{direct_instruction}" if system else direct_instruction
    if prompt_format == "sft":
        parts = []
        if system:
            parts.append(f"<|system|>\n{system}\n{eos_token}\n")
        for role, content in messages:
            parts.append(f"<|{role}|>\n{content}\n{eos_token}\n")
        assistant_prefix = "<|assistant|>\n"
        # The normalized SFT corpus uses canonical reasoning delimiters. Seed
        # the assistant turn when thinking is explicitly requested so the
        # model cannot skip the reasoning section after following the system
        # instruction only probabilistically.
        if think:
            assistant_prefix += "<think>"
        return "".join(parts) + assistant_prefix
    parts = [f"System: {system}\n" if system else ""]
    for role, content in messages:
        label = "User" if role == "user" else "Assistant"
        parts.append(f"{label}: {content}\n")
    return "".join(parts) + "Assistant:"


def main() -> None:
    parser = argparse.ArgumentParser(description="Chat with a pretrained decoder checkpoint")
    parser.add_argument("--checkpoint", type=Path, default=Path("outputs/checkpoints/decoder-10b-zero-v1/latest.pt"))
    parser.add_argument("--config", type=Path, default=Path("configs/decoder_10b.json"))
    parser.add_argument("--tokenizer", default=None, help="defaults to the tokenizer in the config")
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--mode", choices=["chat", "base"], default="chat",
                        help="chat uses plain User/Assistant turns; base uses the prompt verbatim")
    parser.add_argument("--prompt-format", choices=["plain", "sft"], default="plain",
                        help="chat prompt serialization; use sft for structured SFT checkpoints")
    parser.add_argument(
        "--think", nargs="?", const=True, default=False, type=parse_bool,
        help="enable explicit step-by-step reasoning (accepts true/false)",
    )
    parser.add_argument("--prompt", default=None, help="one-shot prompt; chat mode formats it as a user turn")
    parser.add_argument("--system", default=None)
    parser.add_argument("--max-new-tokens", type=int, default=256)
    parser.add_argument("--temperature", type=float, default=0.8)
    parser.add_argument("--top-p", type=float, default=0.95)
    parser.add_argument("--top-k", type=int, default=0)
    parser.add_argument("--repetition-penalty", type=float, default=1.0)
    args = parser.parse_args()
    if args.max_new_tokens <= 0 or not 0 < args.top_p <= 1 or args.temperature < 0:
        raise ValueError("invalid generation settings")
    device = torch.device(args.device)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError(
            "CUDA was requested but is unavailable in this container. "
            "Check Docker GPU access with nvidia-smi; use --device cpu explicitly "
            "only as a fallback."
        )
    model, config = load_model(args.checkpoint, args.config, device)
    # A published artifact carries its tokenizer beside the safetensors file.
    # Prefer that local copy so inference does not need to contact the Hub.
    local_tokenizer = args.checkpoint.parent
    if args.tokenizer is None and (local_tokenizer / "tokenizer.json").is_file():
        tokenizer_name = str(local_tokenizer)
    else:
        tokenizer_name = args.tokenizer or config.get("tokenizer")
    if not tokenizer_name:
        raise ValueError("provide --tokenizer or set tokenizer in the config")
    tokenizer = load_fast_tokenizer(tokenizer_name)
    if len(tokenizer) != model.config.vocab_size:
        raise ValueError(f"tokenizer vocabulary ({len(tokenizer)}) does not match model ({model.config.vocab_size})")

    if args.mode == "base":
        if args.prompt is None:
            raise ValueError("--prompt is required in --mode base")
        prompt = args.prompt
        if args.think:
            prompt = (
                "Think through the request carefully. Write the reasoning step by step "
                "inside <think>...</think>, then give the final answer after </think>.\n\n"
                + prompt
            )
        # Pretraining shards were encoded without BOS; match that convention
        # at inference instead of injecting an unseen leading special token.
        prompt_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids
        print("completion> ", end="", flush=True)
        generate(model, tokenizer, prompt_ids, args.max_new_tokens, args.temperature, args.top_p,
                 args.top_k, args.repetition_penalty, device)
        return

    if args.prompt is not None:
        history = [("user", args.prompt)]
        prompt = format_chat_prompt(
            history, args.system, args.prompt_format, tokenizer.eos_token or "</s>", args.think
        )
        prompt_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids
        prompt_ids = prompt_ids[:, -model.config.sequence_len:]
        visible_prefix = "<think>" if args.think and args.prompt_format == "sft" else ""
        print(f"assistant> {visible_prefix}", end="", flush=True)
        generate(model, tokenizer, prompt_ids, args.max_new_tokens, args.temperature,
                 args.top_p, args.top_k, args.repetition_penalty, device)
        return

    print("Loaded checkpoint. Type /exit to quit or /clear to reset history.")
    history: list[tuple[str, str]] = []
    while True:
        try:
            user_text = input("user> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not user_text:
            continue
        if user_text in {"/exit", "/quit"}:
            break
        if user_text == "/clear":
            history.clear()
            print("history cleared")
            continue
        history.append(("user", user_text))
        prompt = format_chat_prompt(
            history, args.system, args.prompt_format, tokenizer.eos_token or "</s>", args.think
        )
        prompt_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids
        # Keep the newest conversation turns when history exceeds the context.
        prompt_ids = prompt_ids[:, -model.config.sequence_len:]
        visible_prefix = "<think>" if args.think and args.prompt_format == "sft" else ""
        print(f"assistant> {visible_prefix}", end="", flush=True)
        answer = generate(model, tokenizer, prompt_ids, args.max_new_tokens, args.temperature,
                          args.top_p, args.top_k, args.repetition_penalty, device)
        history.append(("assistant", answer))


if __name__ == "__main__":
    main()