import sys, json, os
sys.path.insert(0, os.environ.get("INDIGO_DIR", r"D:\Documents\GitHub\Indigo"))
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
import torch
from safetensors.torch import load_file
from indigo.common import build_tokenizer
from indigo.model import GPT, GPTConfig

ckpt = os.environ.get("CKPT", r"D:\Documents\GitHub\Indigo\runs\indo\indigo.safetensors")
meta_path = ckpt.replace(".safetensors", "_meta.json")
with open(meta_path, encoding="utf-8") as f:
    meta = json.load(f)
state = load_file(ckpt)
model = GPT(GPTConfig.from_dict(meta["config"]))
model.load_state_dict(state, strict=False)
model.eval()
tokenizer = build_tokenizer(meta.get("tokenizer") or {"type": "char"}, meta.get("vocab"))
block_size = meta["config"]["block_size"]

history_ids = json.loads(os.environ.get("HISTORY", "[]"))
prompt = os.environ.get("PROMPT", "Apa itu Indigo?")

temperature = float(os.environ.get("TEMPERATURE", "0.8"))
top_k = int(os.environ.get("TOP_K", "40"))
top_p = float(os.environ.get("TOP_P", "0.95"))
max_new = int(os.environ.get("MAX_NEW", "200"))

piece = tokenizer.encode("\nAnda: " + prompt + "\nIndigo:")
ctx = (history_ids + piece)[-block_size:]
x = torch.tensor([ctx], dtype=torch.long)
with torch.no_grad():
    out = model.generate(x, max_new, temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=1.15)
full = out[0].tolist()
new_ids = full[len(ctx):]
text = tokenizer.decode(new_ids).strip()
print(json.dumps({"text": text, "ids": full[-block_size:]}))