#!/usr/bin/env python3
"""Test batch naming: exercises naming_core against real Gemma 4.

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

sys.path.insert(0, str(Path(__file__).parent))
from gemma_naming import (
    CORRIDOR_SUBSTYLES,
    DEFAULT_SUBSTYLES,
    _PROMPT_CONFIG,
)
from naming_core import (
    build_batch_prompt,
    parse_batch_response,
    select_distinct,
)

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

# Full system simulations — 5 bodies each, accumulating taken list
SYSTEM_SIMS = [
    {
        "name": "Ran", "corridor": "core",
        "inflection": "English countryside, rural, agricultural settlers",
        "hook": "RAN does not import food.",
        "bodies": [
            ("GJ144b", "ambition"), ("GJ144c", "community"),
            ("GJ144d", "fear"), ("GJ144e", "hope"), ("GJ144e-1", "loss"),
        ],
    },
    {
        "name": "Groombridge", "corridor": "core",
        "inflection": "British colonial settlement era",
        "hook": "GROOMBRIDGE is where the money lives.",
        "bodies": [
            ("GJ380b", "wealth"), ("GJ380c", "pride"),
            ("GJ380d", "industry"), ("GJ380e", "ambition"), ("GJ380f", "fleeting"),
        ],
    },
    {
        "name": "Cairnside", "corridor": "deep_frontier",
        "inflection": "frontier descriptive, geographic features named by surveyors",
        "hook": "CAIRNSIDE is a materials science program running for forty years.",
        "bodies": [
            ("GJ892b", "defiance"), ("GJ892c", "isolation"),
            ("GJ892d", "fear"), ("GJ892e", "hope"), ("GJ892f", "loss"),
        ],
    },
    {
        "name": "Ratnagiri", "corridor": "north_reach",
        "inflection": "South African English settler",
        "hook": "RATNAGIRI has a monopoly on its primary export that no one engineered.",
        "bodies": [
            ("GJ575Ab", "pride"), ("GJ575Ac", "community"),
            ("GJ575Ad", "industry"), ("GJ575Ae", "devotion"), ("GJ575Af", "ambition"),
        ],
    },
]

# Build test list
TESTS = []
for sim in SYSTEM_SIMS:
    for body_id, mood in sim["bodies"]:
        TESTS.append({
            "label": f"{sim['name']} — {body_id} (mood: {mood})",
            "system": sim["name"], "body": body_id, "corridor": sim["corridor"],
            "inflection": sim["inflection"],
            "feature_type": "mountain_range", "count": 8,
            "taken": f"__accumulate_{sim['name']}__",
            "hook": sim["hook"], "mood": mood,
        })


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

    accum: dict[str, list[str]] = {}

    for test in TESTS:
        taken = test["taken"]
        if isinstance(taken, str) and taken.startswith("__accumulate_"):
            key = taken
            taken = list(accum.get(key, []))

        # Primary batch
        prompt = build_batch_prompt(
            feature_type=test["feature_type"],
            inflection=test["inflection"],
            count=test["count"],
            taken=taken,
            prompt_config=_PROMPT_CONFIG,
            system_name=test["system"],
            body_name=test["body"],
            system_hook=test["hook"],
            mood=test.get("mood"),
            ctx_size=CTX_SIZE,
        )

        seed = int(hashlib.sha256(
            f"batch|{test['body']}|{test['feature_type']}".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}"

        candidates = parse_batch_response(raw)
        selected = select_distinct(candidates, test["count"], taken)

        # Refill from adjacent register if short
        if len(selected) < test["count"]:
            shortfall = test["count"] - len(selected)
            refill_taken = taken + selected

            corridor = test.get("corridor", "core")
            substyles = CORRIDOR_SUBSTYLES.get(corridor, DEFAULT_SUBSTYLES)
            primary_idx = next(
                (i for i, s in enumerate(substyles)
                 if s["inflection"] == test["inflection"]),
                0,
            )
            refill_idx = (primary_idx + 1) % len(substyles)
            refill_inflection = substyles[refill_idx]["inflection"]

            refill_prompt = build_batch_prompt(
                feature_type=test["feature_type"],
                inflection=refill_inflection,
                count=shortfall * 3,
                taken=refill_taken,
                prompt_config=_PROMPT_CONFIG,
                system_name=test["system"],
                body_name=test["body"],
                system_hook=test["hook"],
                mood=test.get("mood"),
                ctx_size=CTX_SIZE,
            )
            refill_seed = int(hashlib.sha256(
                f"refill|{test['body']}|{test['feature_type']}".encode()
            ).hexdigest()[:8], 16)

            proc.stdin.write(json.dumps({"prompt": refill_prompt, "seed": refill_seed}) + "\n")
            proc.stdin.flush()
            t1 = time.time()
            refill_line = proc.stdout.readline()
            refill_elapsed = time.time() - t1

            try:
                refill_resp = json.loads(refill_line)
                refill_raw = refill_resp.get("text", "")
            except (json.JSONDecodeError, TypeError):
                refill_raw = ""

            refill_candidates = parse_batch_response(refill_raw)
            extra = select_distinct(refill_candidates, shortfall, refill_taken)
            print(f"  REFILL ({refill_inflection}): {len(refill_candidates)} cand → {len(extra)} new: {extra}")
            selected.extend(extra)

        print(f"  {test['label']}")
        print(f"    {len(candidates)} cand → {len(selected)} selected  ({elapsed:.1f}s)  taken={len(taken)}")
        print(f"    {selected}")

        if isinstance(test["taken"], str) and test["taken"].startswith("__accumulate_"):
            key = test["taken"]
            accum.setdefault(key, []).extend(selected)
            print(f"    [{test['system']}: {len(accum[key])} total]")

    proc.stdin.close()
    try:
        os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
    except Exception:
        pass
    proc.wait()
    print()


if __name__ == "__main__":
    main()