"""Turn the demo form into a /v1/systemone-style request for vev-4b and the answers into label dicts."""
import base64
import io
import os

from PIL import Image

from vev.model import CheckpointEngine

MODEL = os.environ.get("VEV_MODEL", "CountingSheep/vev-4b")
N_QUESTIONS = 4


def load() -> CheckpointEngine:
    return CheckpointEngine(MODEL, dtype="bf16", device="cuda")


def data_url(img: Image.Image) -> str:
    buf = io.BytesIO()
    img.convert("RGB").save(buf, format="PNG")
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("ascii")


def build(rows: list[tuple[str, str]]) -> dict[str, dict]:
    """rows: (question, options). Two or more options separated by ';' make a choice; otherwise yes/no."""
    questions = {}
    for i, (q, opts) in enumerate(rows):
        q = (q or "").strip()
        if not q:
            continue
        labels = [o.strip() for o in (opts or "").split(";") if o.strip()]
        if len(labels) >= 2:
            questions[f"q{i}"] = {"type": "choice", "instructions": q, "criteria": {k: None for k in labels}}
        else:
            questions[f"q{i}"] = {"type": "noul", "instructions": q, "criteria": None}
    return questions


def ask(engine: CheckpointEngine, image: Image.Image, context: str, rows: list[tuple[str, str]]):
    """Returns one {label: probability} dict (or None for an empty row) per row, and the engine time in ms."""
    questions = build(rows)
    if image is None or not questions:
        raise ValueError("Upload an image and write at least one question.")
    state = {"screen": {"image": {"url": data_url(image)}}}
    if (context or "").strip():
        state["context"] = context.strip()
    res = engine.run(state, questions)
    out = []
    for i in range(len(rows)):
        a = res.answers.get(f"q{i}")
        if a is None:
            out.append(None)
        elif a["type"] == "noul":
            out.append({"yes": a["noul"], "no": 1 - a["noul"]})
        else:
            out.append(a["probabilities"])
    return out, res.extensions["timing"]["total_ms"]


EXAMPLE = ["checkout.png", "",
           "Does the screen show an error message?", "",
           "Which checkout step is the user on?", "shipping; payment; review",
           "What should the user do next?", "try another card; wait for the order to ship; nothing, the order went through",
           "", ""]