"""
Local inference script for MicroGPT.
Loads a trained checkpoint and generates text.

Usage:
    # Interactive mode (type prompts)
    python inference.py

    # Generate from a prompt
    python inference.py --prompt "Once upon a time"

    # Use a specific checkpoint
    python inference.py --checkpoint runs_out/micro_gpt_ckpt_step_100000.pt --prompt "Hello"
"""

import os
import sys
import argparse
from pathlib import Path

import torch
import torch.nn.functional as F

# Add current directory for imports
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from models import get_model


def load_model_and_meta(checkpoint_path: str, device: str = 'cpu'):
    """Load a trained model checkpoint and return model + metadata."""
    print(f'Loading checkpoint: {checkpoint_path}')
    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)

    config = ckpt['config']
    meta = ckpt['meta']
    step = ckpt['step']
    best_val_loss = ckpt.get('best_val_loss', 'N/A')

    print(f'  Architecture : {config["arch"]}')
    print(f'  Dataset      : {config["dataset"]}')
    print(f'  Vocab size   : {meta["vocab_size"]}')
    print(f'  Block size   : {config["block_size"]}')
    print(f'  Layers/Heads : {config["n_layer"]}/{config["n_head"]}')
    print(f'  Embed dim    : {config["n_embd"]}')
    print(f'  Parameters   : {ckpt["n_params"]:,}')
    print(f'  Training step: {step}')
    print(f'  Best val loss: {best_val_loss}')

    # Build model
    ModelClass = get_model(config['arch'])
    model = ModelClass(
        vocab_size=meta['vocab_size'],
        block_size=config['block_size'],
        n_layer=config['n_layer'],
        n_head=config['n_head'],
        n_embd=config['n_embd'],
        dropout=config.get('dropout', 0.1),
    )
    model.load_state_dict(ckpt['model'])
    model.to(device)
    model.eval()

    return model, meta, config


def encode_text(text: str, stoi: dict) -> torch.Tensor:
    """Encode a string into token indices using the char-level tokenizer."""
    indices = [stoi.get(c, 0) for c in text]  # 0 = newline (fallback)
    return torch.tensor([indices], dtype=torch.long)


def decode_tokens(tokens: list, itos: dict) -> str:
    """Decode token indices back into a string."""
    return ''.join(itos.get(i, '?') for i in tokens)


@torch.no_grad()
def generate_text(model, prompt: str, meta: dict, max_new_tokens: int = 256,
                  temperature: float = 0.9, top_k: int = 40, device: str = 'cpu'):
    """Generate text continuation from a prompt."""
    stoi = meta['stoi']
    itos = meta['itos']

    # Encode prompt
    prompt_ids = encode_text(prompt, stoi)
    prompt_ids = prompt_ids.to(device)

    # Generate
    output_ids = model.generate(
        prompt_ids,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_k=top_k,
    )

    # Decode
    full_text = decode_tokens(output_ids[0].tolist(), itos)
    return full_text


def interactive_mode(model, meta, device: str):
    """Interactive chat-like generation."""
    print('\n' + '=' * 60)
    print('  MicroGPT Interactive Mode')
    print('  Type your prompt and press Enter.')
    print('  Type "quit", "exit", or "q" to stop.')
    print('=' * 60)

    while True:
        try:
            prompt = input('\nPrompt: ').strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break

        if prompt.lower() in ('quit', 'exit', 'q'):
            break

        if not prompt:
            continue

        print('\nGenerating...\n')
        result = generate_text(model, prompt, meta, device=device)
        print(result)
        print()


def main():
    parser = argparse.ArgumentParser(description='MicroGPT Local Inference')
    parser.add_argument('--checkpoint', type=str,
                        default='runs_out/micro_gpt_ckpt_step_100000.pt',
                        help='Path to model checkpoint')
    parser.add_argument('--prompt', type=str, default=None,
                        help='Text prompt for generation')
    parser.add_argument('--max_new_tokens', type=int, default=256,
                        help='Maximum tokens to generate')
    parser.add_argument('--temperature', type=float, default=0.9,
                        help='Sampling temperature (higher = more random)')
    parser.add_argument('--top_k', type=int, default=40,
                        help='Top-k sampling threshold')
    parser.add_argument('--device', type=str, default='cpu',
                        help='Device to run on (cpu, cuda, mps)')
    args = parser.parse_args()

    # Resolve device
    if args.device == 'cuda' and not torch.cuda.is_available():
        print('CUDA not available, falling back to CPU')
        args.device = 'cpu'
    elif args.device == 'mps' and not hasattr(torch.backends, 'mps') or not torch.backends.mps.is_available():
        print('MPS not available, falling back to CPU')
        args.device = 'cpu'

    # Load model
    model, meta, config = load_model_and_meta(args.checkpoint, args.device)

    if args.prompt:
        # Single generation
        print(f'\nPrompt: {args.prompt}')
        print(f'Generating (temp={args.temperature}, top_k={args.top_k}, max_tokens={args.max_new_tokens})...\n')
        result = generate_text(
            model, args.prompt, meta,
            max_new_tokens=args.max_new_tokens,
            temperature=args.temperature,
            top_k=args.top_k,
            device=args.device,
        )
        print(result)
    else:
        # Interactive mode
        interactive_mode(model, meta, args.device)


if __name__ == '__main__':
    main()