#!/usr/bin/env python3
"""
chat.py: Multi-turn ChatML chat with MetaDiffusion chat-SFT models.

Works with:
  - an exported dir (config.json + model.safetensors + tokenizer/)
  - a training checkpoint (step_*.pt) with --tokenizer

Usage:
    Interactive:
        python3 chat.py --model-path MetaDiffusion-150M-ChatBase/
        python3 chat.py --model-path checkpoints_chat/best.pt \
                        --tokenizer data/tokenizer

    One-shot:
        python3 chat.py --model-path MetaDiffusion-150M-ChatBase/ \
                        --prompt "What is 2+2?" --max-new-tokens 128
"""

import argparse
import json
import math
import sys
from pathlib import Path

import torch
import torch.nn.functional as F
from safetensors.torch import load_file
from transformers import AutoTokenizer

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from model import MetaDiffusionLM, MetaDiffusionConfig  # noqa: E402

MASK_TOKEN_ID = 32000
CHAT_TOKENS = ["<|im_start|>", "<|im_end|>"] + [f"<|r{i}|>" for i in range(1, 8)]
IM_START, IM_END = "<|im_start|>", "<|im_end|>"


def build_config(config_dict):
    valid = {k: v for k, v in config_dict.items()
             if k in MetaDiffusionConfig.__dataclass_fields__}
    config = MetaDiffusionConfig(**valid)
    config.tie_word_embeddings = False
    return config


def load_model(model_path, device):
    path = Path(model_path)
    if path.is_dir():
        with open(path / "config.json") as f:
            config = build_config(json.load(f))
        model = MetaDiffusionLM(config).to(device)
        sd = load_file(path / "model.safetensors")
        sd = {k[len("model."):] if k.startswith("model.") else k: v
              for k, v in sd.items()}
        missing, unexpected = model.load_state_dict(sd, strict=False)
        if missing or unexpected:
            print(f"  Warning: missing={missing[:3]} unexpected={unexpected[:3]}")
    else:
        ckpt = torch.load(model_path, map_location=device, weights_only=False)
        config = build_config(ckpt["config"])
        model = MetaDiffusionLM(config).to(device)
        sd = {k.replace("_orig_mod.", "", 1) if k.startswith("_orig_mod.") else k: v
              for k, v in ckpt["model_state_dict"].items()}
        model.load_state_dict(sd)
    print(f"  Loaded {sum(p.numel() for p in model.parameters())/1e6:.1f}M params, "
          f"vocab={config.mask_vocab_size}")
    return model


def ensure_chat_tokens(tokenizer):
    """Add ChatML + rainbow tokens if missing (base tokenizer case).

    Id 32000 is the diffusion [MASK] id, so a reserved filler token takes it
    first; chat tokens must land at 32001..32009.
    """
    if tokenizer.convert_tokens_to_ids(IM_START) == tokenizer.unk_token_id:
        if len(tokenizer) == 32000:
            tokenizer.add_special_tokens({"additional_special_tokens": ["<|reserved|>"]})
        tokenizer.add_special_tokens({"additional_special_tokens": CHAT_TOKENS})
        assert tokenizer.convert_tokens_to_ids(IM_END) == 32002, \
            "chat token ids wrong (collide with mask id 32000)"
    return tokenizer


def format_messages(messages):
    parts = []
    for m in messages:
        parts.append(f"{IM_START}{m['role']}\n{m['content']}{IM_END}")
    return "\n".join(parts)


def cumulative_unmask_frac(i, N):
    return 0.5 * (1 - math.cos(math.pi * i / N))


@torch.no_grad()
def generate_response(model, tokenizer, prompt_ids, gen_len, num_steps,
                      temperature, repetition_penalty, device, watch=False,
                      stop_on_end=True):
    """Denoise a block of [MASK] tokens after the prompt (inference.py schedule).

    With stop_on_end=True, denoising halts as soon as <|im_end|> or </s> is
    committed in the response region: committed tokens never change, so the
    output is identical, but we skip filling garbage after the terminator.
    """
    model.eval()
    total_len = prompt_ids.shape[1] + gen_len
    x = torch.full((1, total_len), MASK_TOKEN_ID, device=device, dtype=torch.long)
    x[0, : prompt_ids.shape[1]] = prompt_ids
    mask_id = MASK_TOKEN_ID
    im_end_id = tokenizer.convert_tokens_to_ids(IM_END)
    eos_id = tokenizer.eos_token_id
    prompt_len = prompt_ids.shape[1]

    for i in range(num_steps):
        frac_now = cumulative_unmask_frac(i, num_steps)
        frac_next = cumulative_unmask_frac(i + 1, num_steps)

        n_masked = (x == mask_id).sum().item()
        n_total = int((frac_next - frac_now) * gen_len + 0.5)
        if i == num_steps - 1:
            n_unmask = n_masked
        else:
            n_unmask = max(n_total, 1) if n_masked > 0 else 0

        t = 1.0 - frac_now
        logits = model(x, torch.full((1,), t, device=device))

        # Never predict the mask token
        logits[:, :, mask_id] = -1e9

        # Repetition penalty over everything already on the sequence
        if repetition_penalty != 1.0:
            for tok in x[0].unique():
                ti = tok.item()
                logits[0, :, ti] = torch.where(
                    logits[0, :, ti] < 0,
                    logits[0, :, ti] * repetition_penalty,
                    logits[0, :, ti] / repetition_penalty,
                )

        mask_positions = x == mask_id
        mask_logits = logits[mask_positions]
        probs = F.softmax(mask_logits / temperature, dim=-1)
        sampled = torch.multinomial(probs, 1).squeeze(-1)
        mask_flat = mask_positions.nonzero(as_tuple=False)

        if n_unmask < mask_positions.sum():
            # Left-to-right commit: fill the leftmost masked positions first
            # (semi-autoregressive block generation, as in LLaDA-MoE eval).
            # Confidence-based commit lets <|im_end|> win the race at ANY
            # position and commits mid-block tokens before position 0, which
            # produced empty responses and fragment-style output on this model.
            fill_positions = mask_flat[:n_unmask]
            for idx, tok in zip(fill_positions, sampled[:n_unmask]):
                x[idx[0], idx[1]] = tok
        else:
            x[mask_positions] = sampled

        if watch:
            remaining = (x == mask_id).sum().item()
            live_toks = [t for t in x[0, prompt_len:].tolist() if t != mask_id]
            partial = tokenizer.decode(
                cut_response(live_toks, tokenizer), skip_special_tokens=True
            ).strip()
            if len(partial) > 70:
                partial = partial[:70] + "..."
            line = (f"step {i+1:3d}/{num_steps} | t={t:.3f} "
                    f"| masks={remaining:3d} | {partial}")
            if sys.stdout.isatty():
                # live in-place update
                sys.stdout.write("\r" + line[:99].ljust(99))
                sys.stdout.flush()
            elif i % max(1, num_steps // 8) == 0:
                print(line)

        # Stop early once the model commits a terminator in the response
        if stop_on_end and ((x[0, prompt_len:] == im_end_id).any() or
                            (x[0, prompt_len:] == eos_id).any()):
            break

    if watch and sys.stdout.isatty():
        sys.stdout.write("\n")
    return x


def cut_response(tokens, tokenizer):
    """Cut generated token list at <|im_end|> or </s>; drop rainbow/pad tokens."""
    im_end_id = tokenizer.convert_tokens_to_ids(IM_END)
    eos_id = tokenizer.eos_token_id
    rainbow_ids = {tokenizer.convert_tokens_to_ids(f"<|r{i}|>") for i in range(1, 8)}
    out = []
    for t in tokens:
        if t == im_end_id or t == eos_id:
            break
        if t in rainbow_ids or t == tokenizer.pad_token_id:
            continue
        out.append(t)
    return out


def run_turn(model, tokenizer, messages, args, device):
    prompt = format_messages(messages) + f"\n{IM_START}assistant\n"
    prompt_ids = torch.tensor([tokenizer.encode(prompt, add_special_tokens=False)],
                              device=device)
    # Retry on empty responses: small models occasionally commit <|im_end|>
    # as the first token. Bump temperature per attempt
    for attempt in range(3):
        x = generate_response(model, tokenizer, prompt_ids, args.max_new_tokens,
                              args.num_steps,
                              args.temperature * (1 + 0.15 * attempt),
                              args.repetition_penalty, device, watch=args.watch)
        response_tokens = x[0, prompt_ids.shape[1]:].tolist()
        response_tokens = cut_response(response_tokens, tokenizer)
        text = tokenizer.decode(response_tokens, skip_special_tokens=True).strip()
        if text:
            return text
    return "(empty response)"


def main():
    parser = argparse.ArgumentParser(description="MetaDiffusion chat (ChatML)")
    parser.add_argument("--model-path", required=True,
                        help="Exported dir (config+safetensors+tokenizer) or step_*.pt")
    parser.add_argument("--tokenizer", default=None,
                        help="Tokenizer dir (needed when --model-path is a step_*.pt)")
    parser.add_argument("--prompt", default=None, help="One-shot prompt (else REPL)")
    parser.add_argument("--system", default="You are a helpful assistant.",
                        help="System prompt for the REPL")
    parser.add_argument("--max-new-tokens", type=int, default=96,
                        help="Response block size)")
    parser.add_argument("--num-steps", type=int, default=128, help="Denoising steps")
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--repetition-penalty", type=float, default=1.5,
                        help="Small diffusion models loop without a strong penalty")
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--watch", action="store_true", help="Show denoising progress")
    args = parser.parse_args()

    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    print(f"[*] Loading model from {args.model_path}")
    model = load_model(args.model_path, device)

    model_path = Path(args.model_path)
    tok_path = args.tokenizer
    if tok_path is None:
        if model_path.is_dir():
            cand = model_path / "tokenizer"
            if not cand.exists() and (model_path / "tokenizer.json").exists():
                cand = model_path  # exported dirs keep the tokenizer at root
            tok_path = str(cand)
    if tok_path is None or not Path(tok_path).exists():
        raise Exception("No tokenizer")
    tokenizer = AutoTokenizer.from_pretrained(str(tok_path))
    tokenizer = ensure_chat_tokens(tokenizer)
    print(f"[*] Tokenizer: {tok_path} (vocab {len(tokenizer)})")

    if args.prompt:
        messages = [{"role": "user", "content": args.prompt}]
        text = run_turn(model, tokenizer, messages, args, device)
        print(f"\nUser: {args.prompt}\nAssistant: {text}\n")
        return

    print("\nMetaDiffusion chat, type 'exit', 'quit' or Ctrl-D to leave.\n")
    messages = [{"role": "system", "content": args.system}]
    while True:
        try:
            user_input = input("You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if user_input.lower() in ("exit", "quit"):
            break
        if not user_input:
            continue
        messages.append({"role": "user", "content": user_input})
        text = run_turn(model, tokenizer, messages, args, device)
        print(f"Assistant: {text}\n")
        messages.append({"role": "assistant", "content": text})


if __name__ == "__main__":
    main()