"""Try Walker -- load the Walker 1B model and generate text.

A minimal, self-contained entry point for users. It reuses the model definition
inside train_walker_1b_continue.py, so put this file next to that trainer (or
pass --trainer). The architecture is inferred from the saved checkpoint, so any
Walker checkpoint works (ongoing training or best).

Examples:
    python try_walker.py
    python try_walker.py --prompt "The migration of the Arctic tern"
    python try_walker.py --ckpt checkpoints/ckpt_ongoing.pt --greedy
    python try_walker.py --device cpu --max-new 32
"""

import argparse
import importlib.util
import os
import sys
from pathlib import Path

import torch


def _first_existing(paths):
    for p in paths:
        if not p:
            continue
        cand = Path(p).expanduser()
        if cand.exists():
            return str(cand.resolve())
    return None


def locate_trainer(explicit):
    here = Path(__file__).resolve().parent
    cwd = Path.cwd()
    found = _first_existing([
        explicit,
        os.environ.get("TB_TRAINER"),
        here / "train_walker_1b_continue.py",
        cwd / "train_walker_1b_continue.py",
        here / "walker_space" / "train_walker_1b_continue.py",
        here / "CPT" / "train_walker_1b_continue.py",
        here / "Experiment" / "train_walker_1b_continue.py",
    ])
    if not found:
        sys.exit(
            "Could not find train_walker_1b_continue.py.\n"
            "Put this script next to it, or pass --trainer PATH."
        )
    return found


def locate_ckpt(explicit):
    here = Path(__file__).resolve().parent
    cwd = Path.cwd()
    found = _first_existing([
        explicit,
        os.environ.get("TB_CKPT_PATH"),
        os.environ.get("TB_CKPT"),
        "checkpoints/ckpt_ongoing.pt",
        here / "checkpoints" / "ckpt_ongoing.pt",
        cwd / "checkpoints" / "ckpt_ongoing.pt",
    ])
    if not found:
        sys.exit(
            "Could not find a checkpoint.\n"
            "Pass --ckpt PATH, for example --ckpt checkpoints/ckpt_ongoing.pt"
        )
    return found


def load_trainer(path):
    spec = importlib.util.spec_from_file_location("walker_trainer", path)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


def pick_device(requested):
    if requested:
        return requested
    return "cuda" if torch.cuda.is_available() else "cpu"


def build(mod, ckpt, device, tokenizer_name):
    print(f"Loading Walker from {ckpt} on {device} ...", flush=True)
    model, tok, cfg = mod.build_model_for_inference(
        ckpt, device=device, tokenizer_name=tokenizer_name
    )
    mod.DEVICE = device
    model.eval()
    print(f"Loaded: D={cfg.D} L={cfg.L} heads={cfg.n_head} vocab={cfg.vocab} "
          f"max_t={cfg.max_t}", flush=True)
    return model, tok, cfg


@torch.inference_mode()
def generate_text(mod, model, tok, prompt, device, max_new, temperature, top_p, rep_penalty, seed):
    torch.manual_seed(int(seed))
    ids = tok(prompt, return_tensors="pt", add_special_tokens=False)["input_ids"]
    max_t = int(getattr(model, "max_t", 4096))
    keep = max_t - int(max_new) - 1
    if ids.shape[1] > keep:
        ids = ids[:, -keep:]
    ids = ids.to(device)
    out = model.generate(
        ids,
        max_new=int(max_new),
        temperature=float(temperature),
        top_p=float(top_p),
        rep_penalty=float(rep_penalty),
    )
    return tok.decode(out[0].tolist(), skip_special_tokens=True)


def parse_args():
    ap = argparse.ArgumentParser(description="Try the Walker 1B model.")
    ap.add_argument("--prompt", default=None, help="Text prompt. Omit for interactive mode.")
    ap.add_argument("--ckpt", default=None, help="Checkpoint path (default: checkpoints/ckpt_ongoing.pt).")
    ap.add_argument("--trainer", default=None, help="Path to train_walker_1b_continue.py.")
    ap.add_argument("--device", default=None, help="cuda or cpu (default: auto).")
    ap.add_argument("--tokenizer", default=None, help="Tokenizer name (default: from checkpoint, else gpt2).")
    ap.add_argument("--max-new", type=int, default=64, help="Number of new tokens to generate.")
    ap.add_argument("--temperature", type=float, default=0.8)
    ap.add_argument("--top-p", type=float, default=0.9)
    ap.add_argument("--rep-penalty", type=float, default=1.1)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--greedy", action="store_true", help="Deterministic decoding (temperature 0).")
    return ap.parse_args()


def main():
    args = parse_args()
    temperature = 0.0 if args.greedy else float(args.temperature)

    trainer = locate_trainer(args.trainer)
    ckpt = locate_ckpt(args.ckpt)
    device = pick_device(args.device)
    if device == "cpu":
        print("Note: running the 1B model on CPU will be very slow.", flush=True)

    mod = load_trainer(trainer)
    model, tok, cfg = build(mod, ckpt, device, args.tokenizer)

    if args.prompt is not None:
        text = generate_text(mod, model, tok, args.prompt, device,
                             args.max_new, temperature, args.top_p, args.rep_penalty, args.seed)
        print("\n" + text + "\n")
        return

    print("Ready. Type a prompt, or :quit to exit.")
    print("Commands: :temp X   :top-p X   :max-new N")
    while True:
        try:
            prompt = input("prompt> ")
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not prompt.strip():
            continue
        if prompt.strip() in (":quit", ":q", ":exit"):
            break
        if prompt.startswith(":temp "):
            temperature = float(prompt.split(None, 1)[1])
            print(f"temperature = {temperature}")
            continue
        if prompt.startswith(":top-p "):
            args.top_p = float(prompt.split(None, 1)[1])
            print(f"top_p = {args.top_p}")
            continue
        if prompt.startswith(":max-new "):
            args.max_new = int(prompt.split(None, 1)[1])
            print(f"max_new = {args.max_new}")
            continue
        text = generate_text(mod, model, tok, prompt, device,
                             args.max_new, temperature, args.top_p, args.rep_penalty, args.seed)
        print("\n" + text + "\n", flush=True)


if __name__ == "__main__":
    main()