import sys
import random
import argparse
import numpy as np
import tensorflow as tf

from indigotf.common import build_tokenizer, load_meta
from indigotf.model import build_gpt, generate

CONFIG_KEYS = ("vocab_size", "block_size", "n_layer", "n_head", "n_embd", "dropout")


def load_model(path):
    from safetensors.numpy import load_file

    meta = load_meta(path)
    if meta.get("backend") != "tensorflow":
        raise SystemExit(
            f"{path} berasal dari backend {meta.get('backend')}, gunakan repo indigo (PyTorch)"
        )
    state = load_file(path)
    by_path = {k.replace("/", "_"): v for k, v in state.items()}
    model = build_gpt(**{k: meta["config"][k] for k in CONFIG_KEYS})
    missing = [v.path for v in model.weights if v.path.replace("/", "_") not in by_path]
    if missing:
        raise SystemExit(f"bobot tidak cocok dengan checkpoint: {missing[:5]}")
    model.set_weights([by_path[v.path.replace("/", "_")] for v in model.weights])
    tokenizer = build_tokenizer(meta.get("tokenizer") or {"type": "char"}, meta.get("vocab"))
    return model, meta, tokenizer


def main():
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    parser = argparse.ArgumentParser(description="Generate teks dari checkpoint Indigo-TF")
    parser.add_argument("--ckpt", default="out/indigo_best.safetensors")
    parser.add_argument("--prompt", default="")
    parser.add_argument("--max-new", type=int, default=300)
    parser.add_argument("--temperature", type=float, default=0.8)
    parser.add_argument("--top-k", type=int, default=40)
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--device", default="auto", choices=["auto", "cpu", "gpu"])
    args = parser.parse_args()

    if args.device == "cpu":
        tf.config.set_visible_devices([], "GPU")
    if args.seed is not None:
        random.seed(args.seed)
        np.random.seed(args.seed)
        tf.random.set_seed(args.seed)

    model, meta, tokenizer = load_model(args.ckpt)

    ids = tokenizer.encode(args.prompt) or [0]
    idx = tf.constant([ids], dtype=tf.int64)
    out = generate(
        model,
        idx,
        args.max_new,
        block_size=meta["config"]["block_size"],
        temperature=args.temperature,
        top_k=args.top_k,
    )
    print(tokenizer.decode(out.numpy()[0].tolist()))


if __name__ == "__main__":
    main()