#!/usr/bin/env python3
"""
test.py - Micro Distillery model chat wrapper.

Loads model.json (a GRPO+VAE-distilled bigram model) and provides
a simple REPL chat interface.

Usage:
    python3 test.py                      # interactive chat
    python3 test.py hello                # one-shot from prompt
    python3 test.py --temperature 1.2    # one-shot with custom temperature
    python3 test.py                      # (interactive also supports --temperature)
"""
import json
import os
import random
import sys

MODEL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "model.json")


def load_model(path=MODEL_PATH):
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def sample_next(transitions, token, rng, temperature=1.0):
    bucket = transitions.get(token)
    if not bucket:
        return None
    tokens = list(bucket.keys())
    inv_t = 1.0 / max(temperature, 1e-3)
    weights = [max(bucket[t], 1e-12) ** inv_t for t in tokens]
    total = sum(weights)
    if total <= 0:
        return None
    r = rng.random() * total
    acc = 0.0
    for t, w in zip(tokens, weights):
        acc += w
        if r <= acc:
            return t
    return tokens[-1]


def generate(model, prompt, max_tokens=32, temperature=0.9):
    transitions = model.get("transitions", {})
    bos = model.get("bos_token", "<|endoftext|>")
    eos = model.get("eos_token", "<|endoftext|>")
    seed = int(model.get("harmonic_seed", 1337)) & 0xFFFFFFFF
    rng = random.Random(seed)

    if prompt and prompt.strip():
        current = prompt.strip()[-1]
        prefix = prompt.strip()
    else:
        current = bos
        prefix = ""

    generated = ""
    for _ in range(max_tokens):
        nxt = sample_next(transitions, current, rng, temperature=temperature)
        if nxt is None or nxt == eos:
            break
        generated += nxt
        current = nxt
    return prefix + generated


def main(temperature=0.9):
    print("=== Micro Distillery Model Chat ===")
    if not os.path.exists(MODEL_PATH):
        print("Error: " + MODEL_PATH + " not found.", file=sys.stderr)
        sys.exit(1)
    model = load_model()
    print("Format:      " + str(model.get("format")))
    print("Vocab size:  " + str(len(model.get("vocab", []))))
    print("Transitions: " + str(len(model.get("transitions", {}))))
    print("Seed:        " + str(model.get("harmonic_seed")))
    print("Temperature: " + str(temperature))
    print("Type /quit to exit.\n")
    while True:
        try:
            prompt = input("You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nBye!")
            break
        if not prompt:
            continue
        if prompt.lower() in ("/quit", "/exit", "quit", "exit"):
            break
        out = generate(model, prompt, temperature=temperature)
        print("Bot: " + out)


if __name__ == "__main__":
    args = sys.argv[1:]
    temperature = 0.9
    prompts = []
    i = 0
    while i < len(args):
        if args[i] == "--temperature":
            if i + 1 < len(args):
                try:
                    temperature = float(args[i + 1])
                except ValueError:
                    print("Warning: invalid --temperature value, using 0.9", file=sys.stderr)
                i += 2
            else:
                print("Warning: --temperature requires a value, using 0.9", file=sys.stderr)
                i += 1
        elif args[i] in ("--help", "-h"):
            print("Usage: python3 test.py [prompt...] [--temperature T]")
            sys.exit(0)
        elif args[i].startswith("--"):
            i += 1
        else:
            prompts.append(args[i])
            i += 1
    if prompts:
        model = load_model()
        print(generate(model, " ".join(prompts), temperature=temperature))
    else:
        main(temperature=temperature)