"""Decision Index interface for JADE."""
from .engine import JadeEngine, UnsupportedInput


class DecisionIndexEngine(JadeEngine):
    name = "jade"
    latency = "In-process vLLM wall time including rendering; model initialization excluded."

    def __call__(self, state, questions):
        from decision_index.engines.base import Unsupported

        try:
            return super().__call__(state, questions)
        except UnsupportedInput as exc:
            raise Unsupported(str(exc)) from exc

    def close(self):
        del self.model

    def warmup(self):
        questions = {
            "warmup": {
                "type": "choice",
                "instructions": "Which color is named?",
                "criteria": {"red": "red", "blue": "blue"},
            }
        }
        for _ in range(2):
            self("The color is red.", questions)

    def runtime(self):
        import torch
        import vllm

        return {
            "torch": torch.__version__,
            "vllm": vllm.__version__,
            "cuda": torch.version.cuda,
            "gpu": torch.cuda.get_device_name(),
        }

    def synchronize(self):
        import torch

        torch.cuda.synchronize()