"""
Laya as an HTTP service.

Laya is a typed-decision model: it takes a STATE (text or JSON) plus typed
questions and answers them in one forward pass with calibrated probabilities —
choice, score (ordinal rubric) and noul (P(true)). It does not write prose, so
it is not a replacement for an LLM; it is for the decisions an app currently
makes with hand-written rules: routing, classifying, verifying.

Backends, auto-detected (LAYA_BACKEND forces one):
  laya      (torch)  — Linux / Windows / any CPU or GPU. Production.
  laya_mlx  (MLX)    — Apple Silicon only. Fast local runs and evals.

Hosting
  • Hugging Face Space, SDK "gradio" — HF runs `python app.py`. Gradio must own
    the server there: its launcher is what HF's runtime and ZeroGPU's startup
    check hook into. Our REST routes are attached to the very app Gradio
    serves, so /health and /v1/* sit on the same URL as the test page.
  • anywhere else — one FastAPI app with Gradio mounted at /ui, served by
    uvicorn (see the Dockerfile).

Endpoints
  GET  /health          what is loaded, and whether the first call is paid for
  POST /v1/decide       generic: {state, questions} -> typed answers
  POST /v1/route        pick one of N candidates (Tally OS template routing)

Auth: every call needs `x-laya-token: <LAYA_TOKEN>` when that variable is set.
A Space URL is public, so set it there.
"""
from __future__ import annotations

import importlib
import json
import os
import time
from contextlib import asynccontextmanager
from typing import Any, Dict, List, Literal, Optional, Union

from fastapi import APIRouter, FastAPI, Header, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

# ── ZeroGPU ────────────────────────────────────────────────────────────────
# ZeroGPU refuses to start a Space unless it finds a function decorated with
# @spaces.GPU ("No @spaces.GPU function detected during startup") — even when
# the work is CPU-only. Laya is a 400M encoder: CPU is fast enough, and one
# warm process beats re-attaching a GPU slice per request. So we register a
# probe and, below, bind it to a real button: the check counts functions it can
# see through Gradio, and a bare module-level definition did not satisfy it.
#
# CUDA is left alone: on ZeroGPU no device is visible until a slice is
# attached, so laya loads on CPU by itself. LAYA_HIDE_CUDA=1 forces the issue
# on a host that exposes a GPU we do not want used.
try:
    import spaces  # noqa: F401 — present on HF Spaces images

    @spaces.GPU(duration=10)
    def zero_gpu_probe() -> str:
        """Never used for real work; it exists so ZeroGPU will start."""
        return "ok"

except Exception:  # noqa: BLE001 — not a Space, or the package is absent

    def zero_gpu_probe() -> str:
        return "no spaces package"


if os.getenv("LAYA_HIDE_CUDA", "0") in ("1", "true", "True"):
    os.environ["CUDA_VISIBLE_DEVICES"] = ""

# ── configuration ───────────────────────────────────────────────────────────
# A ZeroGPU worker re-executes this module; two copies would each load 843 MB
# and the loser of the race for the port kills the Space with "[Errno 98]
# address already in use". The first process to reach __main__ claims the port
# through the environment, which every child inherits.
SERVING_ELSEWHERE = os.environ.get("LAYA_SERVING") == "1"
BACKEND = os.getenv("LAYA_BACKEND", "").strip()
MODEL = os.getenv("LAYA_MODEL", "").strip()
# Empty by default: the right dtype depends on the backend — see agent().
# float16 on a CPU produced NaNs here, which surface as a UNIFORM probability
# distribution and confidence 0.0 rather than as an error (seen 22 Sep 2026 on
# the Space: every choice 0.3333).
DTYPE = os.getenv("LAYA_DTYPE", "").strip()
TOKEN = os.getenv("LAYA_TOKEN", "").strip()
ORIGINS = [o.strip() for o in os.getenv("LAYA_CORS_ORIGINS", "").split(",") if o.strip()]
EAGER = os.getenv("LAYA_EAGER", "1") not in ("0", "false", "False") and not SERVING_ELSEWHERE
# HF sets GRADIO_SERVER_PORT on a gradio Space; PORT elsewhere.
PORT = int(os.getenv("PORT") or os.getenv("GRADIO_SERVER_PORT") or 7860)

DEFAULT_MODELS = {"laya": "convaiinnovations/laya", "laya_mlx": "aac6fef/laya-mlx"}

_state: Dict[str, Any] = {"agent": None, "backend": None, "model": None, "load_ms": None, "error": None}


def _pick_backend() -> str:
    """torch first — it runs everywhere; MLX only where it can."""
    order = [BACKEND] if BACKEND else ["laya", "laya_mlx"]
    for name in order:
        try:
            importlib.import_module(name)
            return name
        except Exception:  # noqa: BLE001 — any import failure means "not this one"
            continue
    raise RuntimeError("neither `laya` nor `laya_mlx` is installed")


def agent():
    """The loaded model, loaded once. Raises 503 with the reason if it cannot."""
    if _state["agent"] is not None:
        return _state["agent"]
    if _state["error"]:
        raise HTTPException(503, f"model not loaded: {_state['error']}")
    t0 = time.time()
    try:
        name = _pick_backend()
        mod = importlib.import_module(name)
        model = MODEL or DEFAULT_MODELS[name]
        # MLX runs fp16 on the GPU happily; torch on a CPU needs fp32.
        dtype = DTYPE or ("float16" if name == "laya_mlx" else "float32")
        # Both packages expose load(model_id, dtype=...); backends that do not
        # take dtype raise TypeError, so fall back to the plain call.
        try:
            a = mod.load(model, dtype=dtype)
        except TypeError:
            a = mod.load(model)
        _state.update(agent=a, backend=name, model=model, dtype=dtype, load_ms=round((time.time() - t0) * 1000))
        return a
    except Exception as e:  # noqa: BLE001
        _state["error"] = f"{type(e).__name__}: {e}"
        raise HTTPException(503, f"model not loaded: {_state['error']}") from e


def warm() -> None:
    """Load at boot so the first real call is fast; a failure is reported by
    /health rather than crashing the container."""
    if not EAGER:
        return
    try:
        agent()
    except Exception:  # noqa: BLE001
        pass


def _auth(token: Optional[str]) -> None:
    if TOKEN and token != TOKEN:
        raise HTTPException(401, "bad or missing x-laya-token")


# ── schemas ─────────────────────────────────────────────────────────────────
class Question(BaseModel):
    type: Literal["choice", "score", "noul"]
    instructions: str
    # Laya scores each option at its own [MASK] token, so `criteria` is a MAP
    # of option -> what that option means. A bare list is accepted and each
    # entry becomes its own description; passing a list to the model itself
    # yields a uniform distribution with confidence 0.0 rather than an error.
    criteria: Optional[Union[Dict[str, str], List[str]]] = None


class DecideIn(BaseModel):
    state: Union[str, Dict[str, Any], List[Any]] = Field(..., description="text or JSON the questions are asked about")
    questions: Dict[str, Question]


class Candidate(BaseModel):
    id: Union[int, str]
    label: str                                  # what the model reads
    hint: Optional[str] = None                  # extra context, not in the label


class RouteIn(BaseModel):
    question: str
    candidates: List[Candidate]
    instructions: str = "Which report answers this question best?"
    top_k: int = 5


def _answers(raw: Any) -> Dict[str, Any]:
    """predict() returns a mapping with `answers`; the raw payload is kept too
    so a caller can read confidence fields this wrapper does not model."""
    if isinstance(raw, dict):
        return raw.get("answers", raw)
    return getattr(raw, "answers", {"value": str(raw)})


def _decide(state: Any, questions: Dict[str, Dict[str, Any]]) -> Dict[str, Any]:
    a = agent()
    t0 = time.time()
    raw = a.predict(state, questions)
    return {"answers": _answers(raw), "raw": raw if isinstance(raw, dict) else None,
            "ms": round((time.time() - t0) * 1000, 1)}


def _route(question: str, candidates: List[Candidate], instructions: str, top_k: int) -> Dict[str, Any]:
    """One choice question whose options are the candidates.

    `criteria` maps each option to what it means — that description is what
    Laya scores against the state, so a candidate's `hint` is worth giving.
    Duplicate labels would collide as dict keys, so they are suffixed.
    """
    options: Dict[str, Candidate] = {}
    for c in candidates:
        key = c.label
        if key in options:
            key = f"{c.label} ({c.id})"
        options[key] = c
    criteria = {k: (c.hint or c.label) for k, c in options.items()}

    out = _decide(question, {"pick": {"type": "choice", "instructions": instructions, "criteria": criteria}})
    ans = out["answers"].get("pick", {}) if isinstance(out["answers"], dict) else {}
    probs = ans.get("probabilities") or ans.get("probs") or {}
    if isinstance(probs, list):                      # some builds return a bare vector
        probs = {k: probs[i] for i, k in enumerate(options) if i < len(probs)}
    ranked = sorted(
        ({"id": c.id, "label": c.label, "p": float(probs.get(k, 0.0))} for k, c in options.items()),
        key=lambda r: r["p"], reverse=True,
    )[: max(1, top_k)]
    chosen = ans.get("choice") or ans.get("value") or ans.get("label")
    picked = options.get(str(chosen))
    top = (next((r for r in ranked if r["id"] == picked.id), None) if picked else None) or (ranked[0] if ranked else None)
    return {"id": top["id"] if top else None, "label": top["label"] if top else None,
            "p": top["p"] if top else None, "confidence": ans.get("confidence"),
            "ranked": ranked, "ms": out["ms"]}


# ── REST routes (attached to whichever server ends up running) ──────────────
api = APIRouter()


@api.get("/health")
def health() -> Dict[str, Any]:
    return {
        "ok": _state["agent"] is not None,
        "backend": _state["backend"],
        "model": _state["model"] or MODEL or "(default)",
        "dtype": _state.get("dtype"),
        "load_ms": _state["load_ms"],
        "error": _state["error"],
        "auth_required": bool(TOKEN),
    }


@api.post("/v1/decide")
def decide(body: DecideIn, x_laya_token: Optional[str] = Header(None)) -> Dict[str, Any]:
    _auth(x_laya_token)
    qs = {}
    for k, v in body.questions.items():
        q = {kk: vv for kk, vv in v.model_dump().items() if vv is not None}
        if isinstance(q.get("criteria"), list):
            q["criteria"] = {c: c for c in q["criteria"]}
        qs[k] = q
    return _decide(body.state, qs)


@api.post("/v1/route")
def route(body: RouteIn, x_laya_token: Optional[str] = Header(None)) -> Dict[str, Any]:
    """One choice question over the candidates, answered with the candidate id.

    This is the Tally OS case: "which template answers this?" — the state is
    the user's question and the criteria are candidate report titles. No ledger
    figures are sent, which is what makes it safe to run off our own infra.
    """
    _auth(x_laya_token)
    if not body.candidates:
        raise HTTPException(400, "candidates must not be empty")
    return _route(body.question, body.candidates, body.instructions, body.top_k)


@asynccontextmanager
async def _lifespan(_app: FastAPI):
    warm()
    yield


# Importable as `app` for `uvicorn app:app` off-Space.
app = FastAPI(title="Laya API", version="1.0.0", docs_url="/docs", lifespan=_lifespan)
app.include_router(api)
if ORIGINS:
    app.add_middleware(
        CORSMiddleware, allow_origins=ORIGINS, allow_methods=["POST", "GET"],
        allow_headers=["content-type", "x-laya-token"], allow_credentials=False,
    )


# ── the page you can try it on ──────────────────────────────────────────────
def build_ui():
    import gradio as gr

    def ui_route(question: str, candidates_text: str):
        rows = []
        for line in (candidates_text or "").splitlines():
            line = line.strip()
            if not line:
                continue
            cid, _, label = line.partition("|")
            rows.append(Candidate(id=cid.strip(), label=(label or cid).strip()))
        if not rows:
            return {"error": "one candidate per line: 71 | Revenue booked this financial year"}
        return _route(question, rows, "Which report answers this question best?", 5)

    def ui_decide(state: str, questions_json: str):
        try:
            qs = json.loads(questions_json or "{}")
        except json.JSONDecodeError as e:
            return {"error": f"questions must be JSON: {e}"}
        return _decide(state, qs)

    default_candidates = "\n".join([
        "71 | Revenue booked this financial year",
        "13 | Today total sales and invoice count",
        "33 | Last 12 months daybook grouped by month",
        "73 | Total purchases between two dates",
    ])
    default_questions = json.dumps(
        {"risk": {"type": "score", "instructions": "How risky is extending credit?",
                  "criteria": ["low", "watch", "high"]}}, indent=2)

    with gr.Blocks(title="Laya API") as demo:
        gr.Markdown("### Laya API — the real interface is `POST /v1/route` and `POST /v1/decide` (see `/docs`)")
        with gr.Tab("route"):
            q = gr.Textbox(label="Question", value="is saal ka total sales kitna hua")
            c = gr.Textbox(label="Candidates (id | label per line)", lines=6, value=default_candidates)
            gr.Button("Route").click(ui_route, [q, c], gr.JSON(label="Result"), api_name="route")
        with gr.Tab("decide"):
            st = gr.Textbox(label="State", lines=3,
                            value="Party asks for 45 days credit; they have 3 bills over 90 days overdue.")
            qj = gr.Textbox(label="Questions (JSON)", lines=8, value=default_questions)
            gr.Button("Decide").click(ui_decide, [st, qj], gr.JSON(label="Result"), api_name="decide")
        with gr.Tab("health"):
            gr.Button("Health").click(lambda: health(), None, gr.JSON(label="Status"), api_name="health")
        with gr.Tab("gpu"):
            # Binding the decorated function to a real event is what makes
            # ZeroGPU's "detected during startup" check pass.
            gr.Button("GPU probe").click(zero_gpu_probe, None, gr.Textbox(label="probe"))
    return demo


if __name__ == "__main__" and not SERVING_ELSEWHERE:
    os.environ["LAYA_SERVING"] = "1"
    demo = build_ui()
    if os.getenv("SPACE_ID"):
        # Gradio 6 defaults to SSR: a NODE proxy owns the public port and
        # forwards to Python on port+1. The proxy only knows Gradio's own
        # routes, so /health and /v1/* were unreachable through it whatever we
        # did to the Python app. ssr_mode=False puts Python back on the port.
        # (launch(show_api=...) was also dropped in 6 — pass only what every
        # version accepts, since the Space's gradio version is not ours to pin.)
        try:
            demo.launch(server_name="0.0.0.0", server_port=PORT, prevent_thread_lock=True, ssr_mode=False)
        except TypeError:       # older gradio: no ssr_mode, no proxy either
            demo.launch(server_name="0.0.0.0", server_port=PORT, prevent_thread_lock=True)
        # Gradio's app ends with a catch-all route, so include_router() here
        # would append ours BEHIND it: GET /health returned the SPA page and
        # POST /v1/route returned 405. Splice them in at the front instead.
        # On ZeroGPU the server runs outside this process, so this is a
        # best-effort extra: the reliable interface on a Space is Gradio's own
        # /gradio_api/call/<name> protocol, wired through api_name= above.
        try:
            _tmp = FastAPI()
            _tmp.include_router(api)
            _ours = [r for r in _tmp.routes if getattr(r, "path", "").startswith(("/health", "/v1"))]
            demo.app.router.routes[0:0] = _ours
        except Exception:  # noqa: BLE001
            pass
        warm()
        try:
            demo.block_thread()
        except AttributeError:            # helper renamed: just stay alive
            while True:
                time.sleep(3600)
    else:
        import gradio as gr
        import uvicorn

        uvicorn.run(gr.mount_gradio_app(app, demo, path="/ui"), host="0.0.0.0", port=PORT)