"""Jegue 0.2: typed decisions (choice / yes-no) with calibrated probabilities.

Two classes, same call:  response, _ = model(state, questions)   or   response = model.ask(state, questions)

  Jegue(mode="direct")      Hugging Face / PyTorch. One forward pass per request, no text generated. mode="routed" adds the learned router: questions it flags get a stock-model
                            rationale (at most 64 tokens, generated with the adapter off) before the answer slot is read.
  JegueVLLM(mode="routed")  vLLM (one instance, CUDA graphs, prefix cache, the LoRA served by vLLM). The fast path; use it for latency. mode: routed | direct | brief.

`state` is a string (the shared context); `questions` is a dict: {"id": {"type": "choice", "instructions": "...", "criteria": {"key": "description", ...}}} or
{"id": {"type": "noul", "instructions": "..."}} (yes/no, returns the probability of yes). The response is {"answers": {"id": {"type": "choice", "choice": key, "probabilities": {...}}}}.
"""
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
ROOT = os.path.dirname(HERE)            # the repository root holds the backbone (config, safetensors, tokenizer, chat template)
WEIGHTS = os.path.join(ROOT, "weights")  # lora/, lora-vllm/, head.pt, router-direct-brief.npz


class Jegue:
    """Hugging Face path. mode: 'direct' (single pass, default) or 'routed'."""

    def __new__(cls, mode="direct", model=ROOT, weights=WEIGHTS, **options):
        if mode == "routed":
            from decision_engine_v6r import RoutedSlotEngine
            eng = RoutedSlotEngine(model=model, run=weights, ckpt=weights, parallel_context="direct", **options)
        elif mode == "direct":
            from decision_engine_v6 import BriefSlotEngine
            eng = BriefSlotEngine(model=model, run=weights, ckpt=weights, context="direct", parallel_context="direct", temperature="none", **options)
        else:
            raise ValueError("mode must be 'direct' or 'routed'")
        eng.ask = lambda state, questions: eng(state, questions)[0]
        return eng


class JegueVLLM:
    """vLLM path (needs `pip install vllm`, one CUDA GPU). mode: 'routed' (default), 'direct' or 'brief'."""

    def __new__(cls, mode="routed", model=ROOT, weights=WEIGHTS, **options):
        from decision_engine_vr import VllmRoutedEngine
        eng = VllmRoutedEngine(model=model, run=weights, ckpt=weights, mode=mode, **options)
        eng.ask = lambda state, questions: eng(state, questions)[0]
        return eng