#!/usr/bin/env python3
"""Quick test: send register-selection prompts to sr-voice for a few
systems and print what Gemma actually picks.

Usage:
    python3 tooling/planet-gen/test_register_selection.py
"""
import json
import subprocess
import sys
import hashlib
import re
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
from gemma_naming import (
    load_wiki_context,
    CORRIDOR_SUBSTYLES,
    DEFAULT_SUBSTYLES,
    palette_for,
    _extract_cultural_lines,
)

BIN = Path.home() / "Projects/settled-reach/binaries/sr-voice-tooling"
MODEL = Path.home() / "Projects/settled-reach/models/gemma-4.gguf"
DISTROBOX = "reach-build"
CTX_SIZE = 1024

TEST_SYSTEMS = [
    ("GJ 411",  "south_reach",   "Lalande — Iberian/Portuguese"),
    ("GJ 884",  "south_reach",   "Matamba — Angolan/Lusophone"),
    ("GJ 506",  "west_reach",    "Dokkum — Dutch/Frisian"),
    ("GJ 581",  "west_reach",    "Rødvik — Nordic"),
    ("GJ 34B",  "east_reach",    "Yongjin — Korean"),
    ("GJ 205",  "east_reach",    "Kurashiki — Japanese"),
    ("GJ 71",   "core",          "Gateway — administrative hub"),
    ("GJ 144",  "core",          "Ran — agricultural"),
]


def build_prompt(system_id: str, corridor: str) -> tuple[str, list[dict]]:
    """Build the register selection prompt. Returns (prompt, substyles)."""
    substyles = CORRIDOR_SUBSTYLES.get(corridor, DEFAULT_SUBSTYLES)
    wiki_text, gttr_text = load_wiki_context(system_id)

    options = []
    for idx, style in enumerate(substyles, 1):
        options.append(f"{idx}. {style['inflection']}")
    option_block = "\n".join(options)

    context_parts = []
    if gttr_text:
        for para in gttr_text.strip().split("\n\n"):
            cleaned = para.replace("#", "").strip()
            if cleaned.startswith("THE DRIFTER") or cleaned.startswith("DRIFTER"):
                continue
            if not cleaned or len(cleaned) < 20:
                continue
            context_parts.append(cleaned[:300])
            break
    if wiki_text:
        cultural = _extract_cultural_lines(wiki_text)
        if cultural:
            context_parts.append(cultural)

    context = "\n".join(context_parts)

    preamble = (
        "Match the star system to the best cultural naming register.\n\n"
        "System: Neustadt — German-heritage industrial town, west corridor, orderly municipal governance.\n"
        "1. German settlement  2. Dutch colonial  3. Nordic  4. Polish/Czech  5. Baltic/Finnish\n"
        "Best: 1\n\n"
        "System: Matsue — Japanese precision manufacturing hub, east corridor.\n"
        "1. Korean  2. Japanese  3. Taiwanese/Hakka  4. Filipino  5. Mixed East Asian\n"
        "Best: 2\n\n"
        "System: "
    )
    tail = f"\n{option_block}\nBest (number only):"

    max_prompt_chars = (CTX_SIZE - 16) * 4
    budget = max_prompt_chars - len(preamble) - len(tail)
    if budget < 100:
        budget = 100
    if len(context) > budget:
        context = context[:budget]

    prompt = f"{preamble}{context}{tail}"
    return prompt, substyles


def main():
    # Show prompt sizes first
    print("Prompt token estimates (rough: chars/4):")
    for system_id, corridor, note in TEST_SYSTEMS:
        prompt, _ = build_prompt(system_id, corridor)
        est_tokens = len(prompt) // 4
        print(f"  {system_id:<8s} ~{est_tokens:>4d} tokens  ({len(prompt)} chars)  {note}")
    print()

    cmd = ["distrobox", "enter", DISTROBOX, "--",
           str(BIN),
           "--model", str(MODEL),
           "--ctx-size", str(CTX_SIZE)]
    print(f"starting sr-voice (ctx_size={CTX_SIZE})...", flush=True)
    proc = subprocess.Popen(
        cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL, text=True, bufsize=1,
    )

    results = []
    for system_id, corridor, note in TEST_SYSTEMS:
        prompt, substyles = build_prompt(system_id, corridor)
        seed = int(hashlib.sha256(
            f"register|{system_id}|0".encode()
        ).hexdigest()[:8], 16)

        req = json.dumps({"prompt": prompt, "seed": seed})
        proc.stdin.write(req + "\n")
        proc.stdin.flush()

        t0 = time.time()
        resp_line = proc.stdout.readline()
        elapsed = time.time() - t0

        try:
            resp = json.loads(resp_line)
            raw = resp.get("text", resp.get("error", ""))
        except (json.JSONDecodeError, TypeError):
            raw = f"ERR:{resp_line!r}"

        # Parse
        digits = re.search(r"\d+", raw.strip() or "")
        if digits:
            try:
                choice = int(digits.group())
                if 1 <= choice <= len(substyles):
                    picked = f"#{choice} {substyles[choice - 1]['inflection']}"
                else:
                    picked = f"OUT OF RANGE ({choice})"
            except ValueError:
                picked = "PARSE FAIL"
        else:
            picked = f"FAIL: {raw[:60]}"

        hash_pal = palette_for(corridor, system_id)
        results.append((system_id, note, raw.strip()[:12], picked,
                        hash_pal["inflection"], elapsed))

    proc.stdin.close()
    proc.wait()

    print()
    print(f"{'System':<8s} {'Raw':<12s}  {'Gemma picked':<48s} {'Hash fallback':<45s} {'Time':>5s}")
    print("-" * 130)
    for system_id, note, raw, picked, hash_pick, elapsed in results:
        print(f"{system_id:<8s} {raw:<12s}  {picked:<48s} {hash_pick:<45s} {elapsed:4.1f}s")
        print(f"  └ {note}")
    print()


if __name__ == "__main__":
    main()