#!/usr/bin/env python3
"""
Interactive streaming text generation REPL for Teensy.

Copyright (c) 2025 Pankaj Doharey
Modified from NanoGPT (Andrej Karpathy).
"""

import os
import sys
import argparse
import codecs
import secrets
from contextlib import nullcontext

import torch
import tiktoken

from model import TeensyConfig, TeensyLM, adapt_nanogpt_weights


def resolve_device(requested):
    """Pick an available device and warn about MPS quality issues."""
    if requested == 'cuda':
        if torch.cuda.is_available():
            return 'cuda'
        fallback = 'mps' if torch.backends.mps.is_available() else 'cpu'
        print(f"CUDA not available; falling back to {fallback}.")
        return fallback
    if requested == 'mps':
        if torch.backends.mps.is_available():
            print("WARNING: MPS inference for this model can produce degraded output")
            print("         (garbled text and stray <|endoftext|> tokens). CPU is recommended.")
            return 'mps'
        print("MPS not available; falling back to cpu.")
        return 'cpu'
    return 'cpu'


def load_teensy(args):
    args.device = resolve_device(args.device)
    device_type = 'cuda' if 'cuda' in args.device else 'mps' if 'mps' in args.device else 'cpu'

    ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[args.dtype]
    ctx = nullcontext() if device_type in ['cpu', 'mps'] else torch.amp.autocast(device_type=device_type, dtype=ptdtype)

    ckpt_path = os.path.join(args.out_dir, 'teensy-0.pt')
    checkpoint = torch.load(ckpt_path, map_location=args.device)
    cfg = TeensyConfig(**checkpoint['model_args'])
    model = TeensyLM(cfg)

    state_dict = checkpoint['model']
    unwanted_prefix = '_orig_mod.'
    for k, v in list(state_dict.items()):
        if k.startswith(unwanted_prefix):
            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
    state_dict = adapt_nanogpt_weights(state_dict)

    model.load_state_dict(state_dict)
    model.eval()
    model.to(args.device)

    enc = tiktoken.get_encoding("gpt2")
    encode = lambda s: enc.encode(s, allowed_special={"<|endoftext|>"})
    token_bytes = enc.decode_single_token_bytes
    eos_token_id = enc.eot_token

    return model, encode, token_bytes, ctx, eos_token_id


def stream_generate(prompt, model, encode, token_bytes, ctx, args, eos_token_id):
    ids = encode(prompt)
    x = torch.tensor(ids, dtype=torch.long, device=args.device)[None, ...]
    utf8_decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
    yield prompt
    with torch.no_grad():
        with ctx:
            for token_id in model.generate_stream(
                x, args.max_new_tokens,
                temperature=args.temperature,
                top_k=args.top_k,
                top_p=args.top_p,
                eos_token_id=eos_token_id,
            ):
                chunk = utf8_decoder.decode(token_bytes(token_id), final=False)
                if chunk:
                    yield chunk
            tail = utf8_decoder.decode(b"", final=True)
            if tail:
                yield tail


def main():
    parser = argparse.ArgumentParser(description='Interactive Teensy text generation')
    parser.add_argument('--out_dir', default='checkpoints', help='Directory containing model checkpoint')
    parser.add_argument('--temperature', type=float, default=0.8, help='Sampling temperature')
    parser.add_argument('--max_new_tokens', type=int, default=500, help='Number of tokens to generate')
    parser.add_argument('--top_k', type=int, default=200, help='Top-k sampling parameter')
    parser.add_argument('--top_p', type=float, default=0.9, help='Nucleus (top-p) sampling parameter')
    parser.add_argument('--device', default='cpu', help='Device to run on (cpu/cuda/mps; cpu recommended)')
    parser.add_argument('--dtype', default='float16', help='Data type for model')
    parser.add_argument('--seed', type=int, default=None, help='Random seed (omit for non-deterministic sampling)')

    args = parser.parse_args()

    seed = args.seed if args.seed is not None else secrets.randbelow(2**32)
    torch.manual_seed(seed)
    if args.device == 'cuda' and torch.cuda.is_available():
        torch.cuda.manual_seed(seed)

    print(f"Loading model from {args.out_dir}...")
    model, encode, token_bytes, ctx, eos_token_id = load_teensy(args)
    print(f"Model loaded! Running on {args.device}")
    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"number of parameters: {n_params/1e6:.2f}M")
    print("No meta.pkl found, assuming GPT-2 encodings...")
    print("\nEnter your prompts (Ctrl+C or type 'exit' to quit):")

    while True:
        try:
            prompt = input("\n> ")
            if prompt.lower() == 'exit':
                break
            sys.stdout.write("\n")
            sys.stdout.flush()
            for chunk in stream_generate(prompt, model, encode, token_bytes, ctx, args, eos_token_id):
                sys.stdout.write(chunk)
                sys.stdout.flush()
            sys.stdout.write("\n")
            sys.stdout.flush()
        except KeyboardInterrupt:
            sys.stdout.write("\nExiting...\n")
            sys.stdout.flush()
            break
        except Exception as e:
            print(f"Error: {e}")


if __name__ == "__main__":
    main()