"""Generate text from a trained checkpoint."""
import argparse
import torch

from model import GPT, GPTConfig
from tokenizer import load_tokenizer


def get_device():
    if torch.backends.mps.is_available():
        return torch.device("mps")
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def load_model(checkpoint_path, device):
    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
    config = GPTConfig(**ckpt["config"])
    model = GPT(config).to(device)
    model.load_state_dict(ckpt["model_state"])
    model.eval()
    return model


def alpaca_prompt(instruction, input_text=""):
    """Format a prompt in Alpaca instruction style (for models trained on Alpaca)."""
    if input_text.strip():
        return f"### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n\n### Response:\n"
    return f"### Instruction:\n{instruction}\n\n### Response:\n"


def generate_text(model, tokenizer, prompt, max_new_tokens=200, temperature=1.0, top_k=40, device="cpu"):
    encoded = tokenizer.encode(prompt)
    if not encoded:
        encoded = [0]
    idx = torch.tensor([encoded], dtype=torch.long, device=device)
    with torch.no_grad():
        out = model.generate(idx, max_new_tokens=max_new_tokens, temperature=temperature, top_k=top_k)
    return tokenizer.decode(out[0].tolist())


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", default="checkpoints/best_model.pt")
    parser.add_argument("--tokenizer", default="tokenizer.json")
    parser.add_argument("--prompt", default="To be or not to be")
    parser.add_argument("--instruction", default=None,
                        help="Use Alpaca-style prompt. Overrides --prompt.")
    parser.add_argument("--input", default="", help="Optional input for Alpaca prompt")
    parser.add_argument("--max_new_tokens", type=int, default=300)
    parser.add_argument("--temperature", type=float, default=0.8)
    parser.add_argument("--top_k", type=int, default=40)
    args = parser.parse_args()

    device = get_device()
    print(f"Device: {device}")

    tokenizer = load_tokenizer(args.tokenizer)
    model = load_model(args.checkpoint, device)
    print(f"Model loaded ({model.num_params():,} params)\n")

    if args.instruction:
        prompt = alpaca_prompt(args.instruction, args.input)
        print(f"Prompt:\n{prompt}")
    else:
        prompt = args.prompt

    result = generate_text(
        model, tokenizer, prompt,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        top_k=args.top_k,
        device=device,
    )
    print(result)