#!/usr/bin/env python3
"""
Interactive chat with the 1B Transformer.
Runs in an infinite conversation loop from the terminal.

Usage:
  python chat.py                                        # auto-find latest checkpoint
  python chat.py /jfs/deepak-kumar/checkpoints/step_19000.pt  # specific checkpoint
"""

import sys
import os
import glob
import time
import torch
import torch.nn.functional as F
import readline  # enables arrow keys and history in input()

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model.config import ModelConfig
from model.transformer import Transformer
from model.data import get_tokenizer


def find_latest_checkpoint():
    """Look for DPO > SFT > pretrained checkpoint."""
    dpo_dir = "/jfs/deepak-kumar/checkpoints_dpo"
    sft_dir = "/jfs/deepak-kumar/checkpoints_sft"
    pt_dir = "/jfs/deepak-kumar/checkpoints"

    # Prefer DPO final
    dpo_final = os.path.join(dpo_dir, "dpo_final.pt")
    if os.path.exists(dpo_final):
        return dpo_final, True

    dpo_files = glob.glob(os.path.join(dpo_dir, "dpo_step_*.pt"))
    if dpo_files:
        return max(dpo_files, key=lambda f: int(f.split("dpo_step_")[1].split(".")[0])), True

    # Then SFT
    sft_final = os.path.join(sft_dir, "sft_final.pt")
    if os.path.exists(sft_final):
        return sft_final, True

    sft_files = glob.glob(os.path.join(sft_dir, "sft_step_*.pt"))
    if sft_files:
        return max(sft_files, key=lambda f: int(f.split("sft_step_")[1].split(".")[0])), True

    # Fall back to pretrained
    pt_files = glob.glob(os.path.join(pt_dir, "step_*.pt"))
    if pt_files:
        return max(pt_files, key=lambda f: int(os.path.basename(f).split("_")[1].split(".")[0])), False

    return None, False


def load_model(checkpoint_path, tokenizer, device="cuda:0"):
    config = ModelConfig()
    model = Transformer(config)
    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)

    # Handle expanded vocab from SFT
    saved_vocab = ckpt.get("vocab_size", config.vocab_size)
    if saved_vocab > config.vocab_size:
        config.vocab_size = saved_vocab
        model = Transformer(config)

    model.load_state_dict(ckpt["model"])
    model = model.to(device).bfloat16().eval()
    step = ckpt.get("step", "?")
    loss = ckpt.get("loss", "?")
    del ckpt
    torch.cuda.empty_cache()
    return model, config, step, loss


@torch.no_grad()
def generate_stream(model, tokenizer, prompt, max_new_tokens=512,
                    temperature=0.8, top_k=50, top_p=0.9,
                    repetition_penalty=1.15, device="cuda:0",
                    stop_token_ids=None):
    """Generate tokens one at a time, yielding each for streaming output."""
    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
    generated_ids = []
    prev_decoded_len = 0

    if stop_token_ids is None:
        stop_token_ids = set()
    else:
        stop_token_ids = set(stop_token_ids)
    stop_token_ids.add(tokenizer.eos_token_id)

    for _ in range(max_new_tokens):
        if input_ids.shape[1] >= model.config.max_seq_len:
            break

        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            logits, _ = model(input_ids)

        logits = logits[:, -1, :]

        if repetition_penalty != 1.0 and generated_ids:
            prev_tokens = torch.tensor(generated_ids, device=device).unique()
            for token_id in prev_tokens:
                if logits[0, token_id] > 0:
                    logits[0, token_id] /= repetition_penalty
                else:
                    logits[0, token_id] *= repetition_penalty

        logits = logits / temperature

        if top_k > 0:
            topk_vals, _ = torch.topk(logits, top_k)
            logits[logits < topk_vals[:, -1:]] = float("-inf")

        if top_p < 1.0:
            sorted_logits, sorted_idx = torch.sort(logits, descending=True)
            cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
            mask = cum_probs - F.softmax(sorted_logits, dim=-1) >= top_p
            sorted_logits[mask] = float("-inf")
            logits = sorted_logits.scatter(1, sorted_idx, sorted_logits)

        probs = F.softmax(logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)
        token_id = next_token.item()

        # Stop on any stop token (EOS, <|end|>, <|user|>)
        if token_id in stop_token_ids:
            break

        generated_ids.append(token_id)
        input_ids = torch.cat([input_ids, next_token], dim=1)

        full_decoded = tokenizer.decode(generated_ids, skip_special_tokens=True)
        new_text = full_decoded[prev_decoded_len:]
        prev_decoded_len = len(full_decoded)
        yield new_text

    return


def print_banner(step, loss, device):
    print("\033[1;36m")  # cyan bold
    print("=" * 60)
    print("   1B TRANSFORMER — Interactive Chat")
    print("=" * 60)
    print(f"\033[0m  Checkpoint : step {step}")
    print(f"  Loss       : {loss}")
    print(f"  Device     : {device}")
    print(f"  Parameters : 1.106B")
    print()
    print("  \033[90mCommands:\033[0m")
    print("    \033[33m/quit\033[0m      — exit")
    print("    \033[33m/clear\033[0m     — clear conversation context")
    print("    \033[33m/temp N\033[0m    — set temperature (default 0.8)")
    print("    \033[33m/tokens N\033[0m  — set max tokens (default 512)")
    print("    \033[33m/topp N\033[0m    — set top-p (default 0.9)")
    print("    \033[33m/topk N\033[0m    — set top-k (default 50)")
    print("    \033[33m/rep N\033[0m     — set repetition penalty (default 1.15)")
    print()
    print("\033[90m" + "─" * 60 + "\033[0m")


def main():
    device = "cuda:0"

    is_sft = False
    if len(sys.argv) > 1:
        checkpoint = sys.argv[1]
        is_sft = "sft" in checkpoint.lower()
    else:
        result = find_latest_checkpoint()
        if result[0] is None:
            print("No checkpoint found!")
            sys.exit(1)
        checkpoint, is_sft = result

    tokenizer = get_tokenizer()

    # Add chat tokens for SFT models
    if is_sft:
        special_tokens = ["<|user|>", "<|assistant|>", "<|end|>"]
        vocab = tokenizer.get_vocab()
        new_tokens = [t for t in special_tokens if t not in vocab]
        if new_tokens:
            tokenizer.add_tokens(new_tokens, special_tokens=True)

    print(f"\n  Loading model from {checkpoint}...")
    print(f"  Mode: {'SFT (chat)' if is_sft else 'Base (completion)'}")
    model, config, step, loss = load_model(checkpoint, tokenizer, device)
    print(f"  Model loaded!\n")

    print_banner(step, loss, device)
    if is_sft:
        print("  \033[1;32mSFT mode: The model will respond as a chat assistant.\033[0m\n")

    # Settings
    temperature = 0.7 if is_sft else 0.8
    max_tokens = 512
    top_p = 0.9
    top_k = 50
    rep_penalty = 1.15
    context = ""

    # Chat template tokens for SFT
    USER_START = "<|user|>\n"
    ASST_START = "<|assistant|>\n"
    TURN_END = "\n<|end|>\n"

    # Build stop token IDs for generation
    sft_stop_ids = []
    if is_sft:
        vocab = tokenizer.get_vocab()
        for tok_str in ["<|end|>", "<|user|>"]:
            if tok_str in vocab:
                sft_stop_ids.append(vocab[tok_str])

    while True:
        try:
            user_input = input("\n\033[1;32mYou:\033[0m ").strip()
        except (KeyboardInterrupt, EOFError):
            print("\n\nGoodbye!")
            break

        if not user_input:
            continue

        # Handle commands
        if user_input.startswith("/"):
            cmd = user_input.lower().split()
            if cmd[0] == "/quit":
                print("Goodbye!")
                break
            elif cmd[0] == "/clear":
                context = ""
                print("\033[90m  [Context cleared]\033[0m")
                continue
            elif cmd[0] == "/temp" and len(cmd) > 1:
                temperature = float(cmd[1])
                print(f"\033[90m  [Temperature set to {temperature}]\033[0m")
                continue
            elif cmd[0] == "/tokens" and len(cmd) > 1:
                max_tokens = int(cmd[1])
                print(f"\033[90m  [Max tokens set to {max_tokens}]\033[0m")
                continue
            elif cmd[0] == "/topp" and len(cmd) > 1:
                top_p = float(cmd[1])
                print(f"\033[90m  [Top-p set to {top_p}]\033[0m")
                continue
            elif cmd[0] == "/topk" and len(cmd) > 1:
                top_k = int(cmd[1])
                print(f"\033[90m  [Top-k set to {top_k}]\033[0m")
                continue
            elif cmd[0] == "/rep" and len(cmd) > 1:
                rep_penalty = float(cmd[1])
                print(f"\033[90m  [Repetition penalty set to {rep_penalty}]\033[0m")
                continue
            else:
                print("\033[90m  Unknown command. Try /quit, /clear, /temp, /tokens, /topp, /topk, /rep\033[0m")
                continue

        # Build prompt
        if is_sft:
            prompt = context + USER_START + user_input + TURN_END + ASST_START
        else:
            if context:
                prompt = context + "\n" + user_input
            else:
                prompt = user_input

        # Trim context if too long
        while len(tokenizer.encode(prompt)) > config.max_seq_len - max_tokens:
            if is_sft:
                parts = context.split(TURN_END)
                if len(parts) <= 2:
                    break
                context = TURN_END.join(parts[2:])
                prompt = context + USER_START + user_input + TURN_END + ASST_START
            else:
                lines = prompt.split("\n")
                if len(lines) <= 2:
                    break
                prompt = "\n".join(lines[1:])

        # Generate with streaming
        print("\033[1;34mModel:\033[0m ", end="", flush=True)
        t0 = time.time()
        full_response = ""
        token_count = 0

        for token_text in generate_stream(
            model, tokenizer, prompt,
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            repetition_penalty=rep_penalty,
            device=device,
            stop_token_ids=sft_stop_ids if is_sft else None,
        ):
            print(token_text, end="", flush=True)
            full_response += token_text
            token_count += 1

        elapsed = time.time() - t0
        tps = token_count / max(elapsed, 1e-9)
        print(f"\n\033[90m  [{token_count} tokens, {tps:.1f} tok/s, {elapsed:.1f}s]\033[0m")

        # Append to context for multi-turn
        if is_sft:
            context = (context + USER_START + user_input + TURN_END +
                       ASST_START + full_response.strip() + TURN_END)
        else:
            context = prompt + full_response


if __name__ == "__main__":
    main()