"""
test_concept.py — verify STAGE 2 (monster design -> strict JSON) in isolation.

Sends the same system + JSON prompt the server uses, then checks the reply parses
into a monster spec with the required keys. Run `inspect_api.py` first to confirm
the /chat arg order below.

The server is defensive: it clamps out-of-range numbers and falls back if `type`
is off-list, so the HARD requirement is only that the reply *parses* and has the
keys. Out-of-range values are printed as soft warnings, not failures.

Consumes a little ZeroGPU quota (one call).

Usage:
    HF_TOKEN=hf_xxx python test_concept.py
    HF_TOKEN=hf_xxx python test_concept.py --descriptor "wooden chair"
"""

import os
import re
import sys
import json
import argparse

from gradio_client import Client

SPACE = os.getenv("CONCEPT_SPACE", "huggingface-projects/gemma-4-12b-it")
TOKEN = os.getenv("HF_TOKEN")

TYPES = ["beast", "bug", "aquatic", "flora", "mineral",
         "space", "machina", "structure", "culture", "cuisine"]

SYSTEM = (
    "You are a creature designer for a monster-collection game called Piclets. Given a "
    "real-world object, you invent ONE original collectible creature inspired by it. You "
    "always reply with exactly one JSON object and nothing else — no prose, no markdown, "
    "no code fences."
)


def build_prompt(descriptor: str) -> str:
    return (
        f'Design a Piclet inspired by this object: "{descriptor}".\n\n'
        "Return a JSON object with EXACTLY these keys and nothing else:\n"
        '- "name": 1-2 words, max 20 chars, must not contain the object name.\n'
        f'- "type": exactly one of {TYPES}.\n'
        '- "appearance": 1-3 sentences for an image generator; no object name, no art style.\n'
        '- "description": 1-2 sentences of flavour.\n'
        '- "weight_kg": a number.\n'
        '- "height_m": a number.\n'
        '- "rarity": an integer 1-100.\n\n'
        "Reply with only the JSON object."
    )


def extract_json(text: str) -> dict:
    """Mirrors app.py _extract_json: strip framing/fences, parse first {...}."""
    text = text.replace("**💬 Response:**", "")
    text = re.sub(r"^\s*assistant(final)?\s*", "", text, flags=re.IGNORECASE)
    text = re.sub(r"```(?:json)?", "", text)
    s, e = text.find("{"), text.rfind("}")
    if s != -1 and e > s:
        text = text[s:e + 1]
    return json.loads(text)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--descriptor", default="ceramic coffee mug")
    args = ap.parse_args()

    print(f"[concept] space = {SPACE}")
    print(f"[concept] descriptor = {args.descriptor!r}")
    if not TOKEN:
        print("[concept] WARNING: no HF_TOKEN — tiny anonymous quota; may fail.")

    client = Client(SPACE, hf_token=TOKEN)

    # huggingface-projects/gemma-4-12b-it /chat positional args (verify with
    # inspect_api.py):
    #   (text, files, history, thinking, max_new_tokens, image_token_budget,
    #    system_prompt, temperature, top_p, top_k, repetition_penalty) -> response
    # thinking=False so the reply is clean JSON, not interleaved reasoning.
    result = client.predict(
        build_prompt(args.descriptor),  # text
        None,          # files
        None,          # history
        False,         # thinking
        2000,          # max_new_tokens
        280,           # image_token_budget
        SYSTEM,        # system_prompt
        0.7,           # temperature
        api_name="/chat",
    )

    print("\n[concept] RAW RESULT:")
    print(repr(result)[:2000])

    # gemma returns {"reasoning": "", "content": "<reply>"}; normalize to the reply
    # string before parsing. Mirror app.py generate_concept's extraction.
    if isinstance(result, dict):
        raw = result.get("content") or result.get("text") or result.get("response") or ""
    elif isinstance(result, (list, tuple)) and result:
        raw = result[0]
    else:
        raw = result
    raw = raw if isinstance(raw, str) else str(raw)
    try:
        data = extract_json(raw)
    except Exception as exc:
        print(f"\n[concept] FAIL — reply did not parse as JSON: {exc}")
        print("        Check RAW RESULT and adjust the prompt or _extract_json framing rules.")
        sys.exit(1)

    print("\n[concept] parsed JSON:")
    print(json.dumps(data, indent=2, ensure_ascii=False))

    required = ["name", "type", "appearance", "description", "weight_kg", "height_m", "rarity"]
    missing = [k for k in required if k not in data]

    # Soft checks — the server clamps/falls-back on these, so they only warn.
    soft = {
        "type in categories": str(data.get("type", "")).lower() in TYPES,
        "weight is a number": isinstance(data.get("weight_kg"), (int, float)),
        "height is a number": isinstance(data.get("height_m"), (int, float)),
        "rarity within 1-100": isinstance(data.get("rarity"), (int, float)) and 1 <= float(data.get("rarity", 0)) <= 100,
    }
    for name, good in soft.items():
        print(f"  [{'ok' if good else '~~'}] {name}")
    if missing:
        print(f"  missing required keys: {missing}")

    ok = not missing  # hard requirement: parses + has all keys
    print("[concept] PASS" if ok else "[concept] FAIL — required keys missing; fix the prompt/model.")
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()