"""JEV-9B Decision API — one clean JSON call over System 1.

Inference core (`jev_core.py`) is the model authors' documented `transformers` + `peft`
path, reused verbatim from the official demo Space (Apache-2.0). This file only wraps it
in a single machine-friendly Gradio endpoint (`/decide`) that returns JSON, instead of the
upstream multi-question HTML UI, so a local agent can call it headlessly.

ZeroGPU rules honoured here:
  * `import spaces` MUST precede torch / transformers / any CUDA-touching import
    (jev_core imports torch, so `spaces` is imported before it).
  * models are placed on `cuda` at MODULE level (inside jev_core, at import) — not
    inside the `@spaces.GPU` function.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")

import spaces  # noqa: E402  -- MUST precede torch / transformers / any CUDA-touching import

import json  # noqa: E402

import gradio as gr  # noqa: E402

import jev_core as jc  # noqa: E402


def _clean_options(raw) -> list:
    """Accept a comma/newline-separated string (or a list) -> list of option strings."""
    if isinstance(raw, (list, tuple)):
        items = [str(o).strip() for o in raw]
    else:
        items = [o.strip() for o in str(raw or "").replace("\n", ",").split(",")]
    return [o for o in items if o]


@spaces.GPU(duration=45)
def decide(kind: str, state: str, question: str, options: str = "") -> str:
    """Answer ONE typed System 1 decision; return calibrated probabilities as JSON.

    kind:     "noul" | "choice" | "score"
    state:    the text to reason about (truncated to 1024 tokens, head 60% / tail 40%)
    question: the decision question
    options:  comma-separated options (choice only, 2-16); ignored for noul/score
    """
    kind = (kind or "").strip()
    state = (state or "").strip()
    question = (question or "").strip()

    if kind not in jc.KINDS:
        return json.dumps({"error": f"kind must be one of {list(jc.KINDS)}"})
    if not state:
        return json.dumps({"error": "state is required"})
    if not question:
        return json.dumps({"error": "question is required"})

    opts = _clean_options(options)
    if kind == "choice":
        if not 2 <= len(opts) <= 16:
            return json.dumps({"error": "choice needs 2-16 comma-separated options"})
        if len(set(opts)) != len(opts):
            return json.dumps({"error": "choice options must be unique"})
    else:
        opts = None

    state_used, truncated, n_tok = jc.truncate_state(state)
    dist = jc.decide(kind, state_used, question, opts)
    best = max(dist, key=dist.get)
    return json.dumps(
        {
            "kind": kind,
            "distribution": dist,
            "best": best,
            "confidence": dist[best],
            "state_tokens": n_tok,
            "truncated": truncated,
        }
    )


with gr.Blocks(title="JEV-9B Decision API") as demo:
    gr.Markdown(
        "## JEV-9B System 1 — typed decisions in one prefill pass\n"
        "`noul` (yes/no) · `choice` (2–16 options) · `score` (0–5). "
        "Returns a calibrated probability distribution as JSON."
    )
    with gr.Row():
        kind_in = gr.Dropdown(list(jc.KINDS), value="noul", label="kind")
        state_in = gr.Textbox(label="state", lines=3)
    question_in = gr.Textbox(label="question")
    options_in = gr.Textbox(label="options (comma-separated; choice only)")
    out_json = gr.Code(label="decision JSON", language="json")
    btn = gr.Button("Decide", variant="primary")
    btn.click(
        decide,
        inputs=[kind_in, state_in, question_in, options_in],
        outputs=[out_json],
        api_name="decide",
    )

demo.launch()