src/alien_obfuscator/riddle_generator.py

"""Riddle generator and LLM backend abstraction.

This module hides the details of *how* a riddle is produced (local model,
HF Inference API, or a hard-coded mock) behind a single ``generate``
function.  It also handles prompt construction, JSON schema validation,
and retry logic.
"""

import json
import logging
import os
import random
from abc import ABC, abstractmethod
from typing import Any

from alien_obfuscator.config import (
    HF_API_TIMEOUT,
    LLM_MAX_TOKENS,
    LLM_TEMPERATURE,
    LOCAL_QUANTIZE,
    LOCAL_TIMEOUT,
    MAX_PLAINTEXT_LENGTH,
    MAX_RETRIES,
    MODAL_TIMEOUT,
    MOCK_DISTRACTORS,
    NUM_OPTIONS,
    OPENCODE_GO_TIMEOUT,
    OPENCODE_GO_URL,
    OPENROUTER_TIMEOUT,
    OPENROUTER_URL,
)

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# JSON schema that a riddle response must satisfy
# ---------------------------------------------------------------------------
RIDDLE_SCHEMA: dict[str, Any] = {
    "type": "object",
    "properties": {
        "riddle": {"type": "string"},
        "options": {
            "type": "array",
            "items": {"type": "string"},
            "minItems": NUM_OPTIONS,
            "maxItems": NUM_OPTIONS,
        },
        "correct_index": {"type": "integer", "minimum": 0, "maximum": NUM_OPTIONS - 1},
        "theme": {"type": "string"},
    },
    "required": ["riddle", "options", "correct_index", "theme"],
}


# ---------------------------------------------------------------------------
# Prompt templates
# ---------------------------------------------------------------------------
SYSTEM_PROMPT_TEMPLATE: str = (
    "You are a mischievous alien archaeologist who has spent centuries studying "
    "ancient human texts. You craft riddles in the voice of {theme_description}.\n\n"
    "Generate a riddle whose ANSWER is: {plaintext}\n\n"
    "Rules:\n"
    "- The riddle must be solvable by a human familiar with {theme_name}, "
    "but confusing to anyone without cultural context.\n"
    "- The riddle should be 2–4 sentences, poetic, and contain at least one clever twist.\n"
    "- Generate exactly {num_options} answer options: 1 correct (the PLAINTEXT itself), "
    "{num_distractors} plausible distractors.\n"
    "- Distractors should be thematically adjacent (same domain, same era, similar concepts).\n"
    "- Output ONLY valid JSON.\n"
)

STRICT_JSON_PROMPT: str = (
    "\n\nIMPORTANT: Return ONLY a raw JSON object. No explanations, no reasoning, "
    "no chain-of-thought, no markdown, no code fences. "
    "Output nothing except the JSON object itself.\n"
    'Schema: {"riddle": "string", "options": ["string", "string", "string", "string", "string"], '
    '"correct_index": 0, "theme": "string"}'
)


def _validate_riddle_json(data: dict[str, Any]) -> dict[str, Any]:
    """Validate a parsed JSON dict against the riddle schema.

    Parameters
    ----------
    data : dict[str, Any]
        The parsed JSON object from the LLM.

    Returns
    -------
    dict[str, Any]
        The validated data dict (unchanged).

    Raises
    ------
    ValueError
        If any required field is missing or has the wrong type / length.
    """
    required = RIDDLE_SCHEMA["required"]
    for key in required:
        if key not in data:
            logger.warning("Missing required field: %s", key)
            raise ValueError(f"Missing required field: {key}")

    if not isinstance(data["riddle"], str) or not data["riddle"].strip():
        logger.warning("Field 'riddle' is empty or not a string")
        raise ValueError("Field 'riddle' must be a non-empty string.")

    opts = data["options"]
    if not isinstance(opts, list) or len(opts) != NUM_OPTIONS:
        logger.warning(
            "Field 'options' has wrong type or length: type=%s, length=%d",
            type(opts).__name__,
            len(opts) if isinstance(opts, list) else -1,
        )
        raise ValueError(f"Field 'options' must be a list of exactly {NUM_OPTIONS} strings.")
    for o in opts:
        if not isinstance(o, str):
            logger.warning("Non-string option found: %r", o)
            raise ValueError("Every item in 'options' must be a string.")

    ci = data["correct_index"]
    if not isinstance(ci, int) or ci < 0 or ci >= NUM_OPTIONS:
        logger.warning("Invalid correct_index: %r", ci)
        raise ValueError(f"Field 'correct_index' must be an integer between 0 and {NUM_OPTIONS - 1}.")

    if not isinstance(data.get("theme", ""), str):
        logger.warning("Field 'theme' is not a string: %r", data.get("theme"))
        raise ValueError("Field 'theme' must be a string.")

    return data


# ---------------------------------------------------------------------------
# LLM backends
# ---------------------------------------------------------------------------
class LLMBackend(ABC):
    """Abstract interface for an LLM inference backend."""

    @abstractmethod
    def generate(self, prompt: str) -> str:
        """Send ``prompt`` to the model and return the raw text response.

        Parameters
        ----------
        prompt : str
            The fully formatted prompt.

        Returns
        -------
        str
            Raw model output.
        """
        ...


class MockBackend(LLMBackend):
    """Hard-coded backend that returns predictable JSON for testing.

    Useful for offline development, CI, and rapid UI iteration without
    waiting for real model inference.
    """

    def generate(self, prompt: str) -> str:
        """Return a canned riddle JSON based on the plaintext in the prompt.

        Parameters
        ----------
        prompt : str
            Ignored except for extracting the plaintext answer.

        Returns
        -------
        str
            A JSON string matching the riddle schema.
        """
        # Try to extract the plaintext from the prompt line "ANSWER is: X"
        plaintext = "the secret message"
        for line in prompt.splitlines():
            if "ANSWER is:" in line:
                plaintext = line.split("ANSWER is:", 1)[-1].strip()
                break

        distractors = list(MOCK_DISTRACTORS)
        return json.dumps(
            {
                "riddle": (
                    "I am the thing that humans whisper when the stars are right, "
                    "the phrase that unlocks the hidden door. What am I?"
                ),
                "options": [plaintext] + distractors,
                "correct_index": 0,
                "theme": "mock",
            },
            indent=2,
        )


class HuggingFaceBackend(LLMBackend):
    """Backend that calls the Hugging Face Inference API (serverless).

    Parameters
    ----------
    model_id : str
        Hugging Face model identifier (e.g. ``"google/gemma-4-31b-it"``).
    api_token : str | None
        Hugging Face API token. If ``None``, the token is read from the
        ``HF_TOKEN`` environment variable.
    """

    def __init__(self, model_id: str, api_token: str | None = None) -> None:
        self.model_id = model_id
        self.api_token = api_token

    def generate(self, prompt: str) -> str:
        """Call the Hugging Face Inference API and return the generated text.

        Parameters
        ----------
        prompt : str
            The prompt to send.

        Returns
        -------
        str
            Raw model output.

        Raises
        ------
        RuntimeError
            If the API request fails or returns an error.
        """
        import os

        import requests

        token = self.api_token or os.environ.get("HF_TOKEN")
        if not token:
            logger.error("HF_TOKEN not found in environment")
            raise RuntimeError("HF_TOKEN not provided and not found in environment.")

        url = f"https://api-inference.huggingface.co/models/{self.model_id}"
        headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
        payload = {
            "inputs": prompt,
            "parameters": {
                "max_new_tokens": LLM_MAX_TOKENS,
                "temperature": LLM_TEMPERATURE,
                "return_full_text": False,
            },
        }

        response = requests.post(url, headers=headers, json=payload, timeout=HF_API_TIMEOUT)
        if response.status_code != 200:
            logger.error("HF API error %d: %s", response.status_code, response.text)
            raise RuntimeError(f"HF API error {response.status_code}: {response.text}")

        data = response.json()
        if isinstance(data, list) and len(data) > 0:
            return data[0].get("generated_text", "")
        logger.warning("HF API returned unexpected format: %s", str(data)[:200])
        return str(data)


class OpenAICompatibleBackend(LLMBackend):
    """Generic backend for any OpenAI-compatible chat completions API.

    Parameters
    ----------
    model_id : str
        Model identifier (e.g. ``"google/gemma-4-31b-it"``).
    api_key : str | None
        API key. If ``None``, read from ``key_env_var`` env variable.
    api_url : str
        The chat completions endpoint URL.
    key_env_var : str
        Environment variable name to look for the API key.
    provider_name : str
        Human-readable provider name for error messages.
    extra_headers : dict
        Additional HTTP headers to send with each request.
    """

    def __init__(
        self,
        model_id: str,
        api_key: str | None = None,
        api_url: str = "",
        key_env_var: str = "",
        provider_name: str = "API",
        extra_headers: dict | None = None,
        timeout: int = 120,
    ) -> None:
        self.model_id = model_id
        self.api_key = api_key
        self._api_url = api_url
        self._key_env_var = key_env_var
        self._provider_name = provider_name
        self._extra_headers = extra_headers or {}
        self._timeout = timeout

    def generate(self, prompt: str) -> str:
        """Call the chat completions API and return the generated text.

        Parameters
        ----------
        prompt : str
            The prompt to send.

        Returns
        -------
        str
            Raw model output.

        Raises
        ------
        RuntimeError
            If the API request fails or returns an error.
        """
        import os

        import requests

        key = self.api_key or os.environ.get(self._key_env_var)
        if not key:
            logger.error("%s API key not found in environment var %s", self._provider_name, self._key_env_var)
            raise RuntimeError(
                f"{self._provider_name} API key not provided and {self._key_env_var} not found in environment."
            )

        headers = {
            "Authorization": f"Bearer {key}",
            "Content-Type": "application/json",
            **self._extra_headers,
        }
        payload = {
            "model": self.model_id,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": LLM_MAX_TOKENS,
            "temperature": LLM_TEMPERATURE,
        }

        response = requests.post(self._api_url, headers=headers, json=payload, timeout=self._timeout)
        if response.status_code != 200:
            logger.error("%s API error %d: %s", self._provider_name, response.status_code, response.text)
            raise RuntimeError(f"{self._provider_name} API error {response.status_code}: {response.text}")

        data = response.json()
        choices = data.get("choices", [])
        if not choices:
            logger.error("%s returned no choices. Response: %s", self._provider_name, str(data)[:500])
            raise RuntimeError(f"{self._provider_name} returned no choices.")
        msg = choices[0].get("message", {})
        content = msg.get("content", "")
        # Reasoning models (e.g. DeepSeek) sometimes put output in
        # reasoning_content instead of content.
        if not content:
            content = msg.get("reasoning_content", "")
            if content:
                logger.warning(
                    "%s returned content in reasoning_content field (model=%s)",
                    self._provider_name,
                    data.get("model", "unknown"),
                )
        if not content:
            logger.error(
                "%s returned empty content. Full response: %s",
                self._provider_name,
                str(data)[:500],
            )
            raise RuntimeError(f"{self._provider_name} returned empty content.")
        return content


class OpenRouterBackend(OpenAICompatibleBackend):
    """Backend for OpenRouter's chat completions API."""

    def __init__(self, model_id: str, api_key: str | None = None) -> None:
        super().__init__(
            model_id=model_id,
            api_key=api_key,
            api_url=OPENROUTER_URL,
            key_env_var="OPENROUTER_API_KEY",
            provider_name="OpenRouter",
            extra_headers={
                "HTTP-Referer": "https://github.com/koala/alien-obfuscator",
                "X-Title": "Alien Obfuscator",
            },
            timeout=OPENROUTER_TIMEOUT,
        )


class OpenCodeGoBackend(OpenAICompatibleBackend):
    """Backend for OpenCode Go chat completions API."""

    def __init__(self, model_id: str, api_key: str | None = None) -> None:
        super().__init__(
            model_id=model_id,
            api_key=api_key,
            api_url=OPENCODE_GO_URL,
            key_env_var="OPENCODE_GO_API_KEY",
            provider_name="OpenCode Go",
            timeout=OPENCODE_GO_TIMEOUT,
        )


class ModalBackend(OpenAICompatibleBackend):
    """Backend for a Modal-deployed vLLM server (OpenAI-compatible).

    Connects to a pre-deployed Modal vLLM instance serving an LLM.
    The Modal API URL is public by default — no authentication is needed
    beyond the URL itself.

    Parameters
    ----------
    model_id : str
        Model identifier served by the Modal endpoint
        (e.g. ``"google/gemma-4-31b-it"``).
    api_url : str | None
        Base URL of the Modal-deployed vLLM server. If ``None``, read from
        the ``MODAL_API_URL`` environment variable.
    """

    def __init__(
        self,
        model_id: str,
        api_url: str | None = None,
    ) -> None:
        import os

        resolved_url = api_url or os.environ.get("MODAL_API_URL", "")
        if not resolved_url:
            raise ValueError("Modal API URL not provided and MODAL_API_URL not set in environment.")

        super().__init__(
            model_id=model_id,
            api_key=None,
            api_url=resolved_url.rstrip("/") + "/v1/chat/completions",
            key_env_var="",
            provider_name="Modal",
            timeout=MODAL_TIMEOUT,
        )

    def generate(self, prompt: str) -> str:
        """Call the Modal-deployed vLLM server and return the generated text.

        Modal web endpoints are public, so no API key is required.

        Parameters
        ----------
        prompt : str
            The prompt to send.

        Returns
        -------
        str
            Raw model output.

        Raises
        ------
        RuntimeError
            If the API request fails or returns an error.
        """
        import requests

        headers = {"Content-Type": "application/json"}
        payload = {
            "model": self.model_id,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": LLM_MAX_TOKENS,
            "temperature": LLM_TEMPERATURE,
        }

        response = requests.post(self._api_url, headers=headers, json=payload, timeout=self._timeout)
        if response.status_code != 200:
            logger.error(
                "%s API error %d: %s",
                self._provider_name,
                response.status_code,
                response.text,
            )
            raise RuntimeError(f"{self._provider_name} API error {response.status_code}: {response.text}")

        data = response.json()
        choices = data.get("choices", [])
        if not choices:
            logger.error(
                "%s returned no choices. Response: %s",
                self._provider_name,
                str(data)[:500],
            )
            raise RuntimeError(f"{self._provider_name} returned no choices.")
        msg = choices[0].get("message", {})
        content = msg.get("content", "")
        if not content:
            content = msg.get("reasoning_content", "")
        if not content:
            logger.error(
                "%s returned empty content. Full response: %s",
                self._provider_name,
                str(data)[:500],
            )
            raise RuntimeError(f"{self._provider_name} returned empty content.")
        return content


class LocalGPU4BitBackend(LLMBackend):
    """Backend that runs the model locally on the host GPU using transformers.

    Loads a Hugging Face causal LM at startup and runs inference directly on
    the local GPU. Supports fp16, bf16, 8-bit, and 4-bit quantization.
    Designed for Hugging Face Spaces with dedicated GPU hardware
    (e.g. RTX Pro 6000 with 96 GB VRAM).

    Parameters
    ----------
    model_id : str
        Hugging Face model identifier (e.g. ``"google/gemma-4-31b-it"``).
    timeout : int
        Maximum seconds allowed for generation. Default is 600.
    quantize : str
        Quantization mode: ``"fp16"``, ``"bf16"``, ``"8bit"``, or ``"4bit"``.
        Default is ``"fp16"``.

    Raises
    ------
    ImportError
        If ``torch`` or ``transformers`` are not installed.
    RuntimeError
        If no CUDA-capable GPU is available or the model fails to load.
    """

    def __init__(
        self,
        model_id: str,
        timeout: int = LOCAL_TIMEOUT,
        quantize: str = LOCAL_QUANTIZE,
    ) -> None:
        import torch

        self._model_id = model_id
        self._timeout = timeout
        self._quantize = quantize
        self._model = None
        self._tokenizer = None
        self._device = "cuda" if torch.cuda.is_available() else "cpu"

        if self._device != "cuda":
            raise RuntimeError(
                "LocalGPU4BitBackend requires a CUDA-capable GPU, "
                f"but torch.cuda.is_available() returned {torch.cuda.is_available()}"
            )

        logger.info(
            "LocalGPU4BitBackend: model=%s quantize=%s device=%s",
            model_id,
            quantize,
            self._device,
        )

        self._load_model()

    @staticmethod
    def is_available() -> bool:
        """Check whether local GPU inference can be used.

        Returns
        -------
        bool
            ``True`` if ``torch`` and ``transformers`` are importable and a
            CUDA-capable GPU is detected.
        """
        try:
            import torch  # noqa: F401
        except ImportError:
            return False
        return torch.cuda.is_available()

    def _load_model(self) -> None:
        """Load tokenizer and model with the configured quantization settings.

        The model is loaded onto the GPU. Quantization is configured via
        ``BitsAndBytesConfig`` for 4/8-bit or via ``torch_dtype`` for fp16/bf16.

        Raises
        ------
        ImportError
            If ``transformers`` is not installed.
        RuntimeError
            If the model fails to load.
        """
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

        logger.info("Loading tokenizer for %s ...", self._model_id)
        tokenizer = AutoTokenizer.from_pretrained(self._model_id)
        if tokenizer.pad_token is None:  # type: ignore
            tokenizer.pad_token = tokenizer.eos_token  # type: ignore

        quantization_config = None
        torch_dtype: torch.dtype | str = torch.float16

        if self._quantize == "4bit":
            quantization_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch.bfloat16,
            )
            torch_dtype = "auto"
        elif self._quantize == "8bit":
            quantization_config = BitsAndBytesConfig(load_in_8bit=True)
            torch_dtype = "auto"
        elif self._quantize in ("bf16", "bfloat16"):
            torch_dtype = torch.bfloat16

        logger.info(
            "Loading model %s with quantize=%s ...",
            self._model_id,
            self._quantize,
        )
        model = AutoModelForCausalLM.from_pretrained(
            self._model_id,
            quantization_config=quantization_config,
            torch_dtype=torch_dtype,
            device_map="auto",
        )
        logger.info("Model loaded successfully on %s", self._device)

        self._tokenizer = tokenizer
        self._model = model

    def generate(self, prompt: str) -> str:
        """Run inference on the local GPU and return generated text.

        Parameters
        ----------
        prompt : str
            The prompt to send.

        Returns
        -------
        str
            Raw model output.

        Raises
        ------
        RuntimeError
            If the model has not been loaded or generation fails.
        """
        if self._model is None or self._tokenizer is None:
            raise RuntimeError("Model not loaded. Call _load_model() first.")

        import torch

        inputs = self._tokenizer(prompt, return_tensors="pt").to(self._device)

        with torch.no_grad():
            outputs = self._model.generate(  # type: ignore
                **inputs,
                max_new_tokens=LLM_MAX_TOKENS,
                temperature=LLM_TEMPERATURE,
                do_sample=True,
                pad_token_id=self._tokenizer.pad_token_id,
            )

        generated = self._tokenizer.decode(outputs[0], skip_special_tokens=True)
        prompt_len = len(self._tokenizer.decode(inputs["input_ids"][0], skip_special_tokens=True))
        if isinstance(generated, str):
            response = generated[prompt_len:].strip()
        else:
            response = str(generated)[prompt_len:].strip()

        return response


class AutoBackend(LLMBackend):
    """Smart backend that tries a configurable sequence of inference backends.

    On each ``generate`` call the wrapper iterates through the ordered
    fallback list, returning the first successful response. The default
    order is **Modal → local GPU → mock**, configurable via the
    ``AUTO_FALLBACK_ORDER`` environment variable or ``config.yaml``.

    Recognised backend names in the fallback order:

    * ``"modal"`` — ``ModalBackend`` (requires ``MODAL_API_URL`` env var)
    * ``"local"`` — ``LocalGPU4BitBackend`` (requires CUDA + torch/transformers)
    * ``"mock"`` — ``MockBackend`` (always available)

    Parameters
    ----------
    model_id : str
        Model identifier used by both Modal and local backends.
    quantize : str
        Quantization mode for local GPU. Default is ``"fp16"``.
    fallback_order : list[str]
        Ordered list of backend names to try. Default is
        ``["modal", "local", "mock"]``.
    """

    # Map backend name → (class_or_none, availability_check)
    _BACKEND_REGISTRY: dict[str, Any] = {
        "modal": (ModalBackend, lambda: bool(os.environ.get("MODAL_API_URL"))),
        "local": (LocalGPU4BitBackend, LocalGPU4BitBackend.is_available),
        "mock": (MockBackend, lambda: True),
    }

    def __init__(
        self,
        model_id: str,
        quantize: str = LOCAL_QUANTIZE,
        fallback_order: list[str] | None = None,
    ) -> None:
        if fallback_order is None:
            fallback_order = ["modal", "local", "mock"]

        self._backends: list[LLMBackend] = []

        for name in fallback_order:
            name = name.strip().lower()
            if name not in self._BACKEND_REGISTRY:
                logger.warning("AutoBackend: unknown backend '%s', skipping", name)
                continue

            backend_cls, availability_fn = self._BACKEND_REGISTRY[name]
            try:
                if not availability_fn():
                    logger.info("AutoBackend: backend '%s' not available", name)
                    continue
            except Exception as exc:
                logger.warning("AutoBackend: availability check for '%s' failed: %s", name, exc)
                continue

            try:
                if name == "modal":
                    backend = backend_cls(model_id)
                elif name == "local":
                    backend = backend_cls(model_id, quantize=quantize)
                elif name == "mock":
                    backend = backend_cls()
                else:
                    backend = backend_cls(model_id)
                self._backends.append(backend)
                logger.info("AutoBackend: registered '%s'", name)
            except Exception as exc:
                logger.warning("AutoBackend: failed to init '%s': %s", name, exc)

        if not self._backends:
            logger.warning("AutoBackend: no backends available, forcing mock")
            self._backends = [MockBackend()]

        logger.info(
            "AutoBackend: fallback chain = %s",
            [type(b).__name__ for b in self._backends],
        )

    def generate(self, prompt: str) -> str:
        """Try each backend in fallback order, return first successful response.

        Parameters
        ----------
        prompt : str
            The prompt to send.

        Returns
        -------
        str
            Raw model output.
        """
        for i, backend in enumerate(self._backends):
            try:
                return backend.generate(prompt)
            except Exception as exc:
                logger.warning(
                    "AutoBackend: backend %d/%d (%s) failed: %s",
                    i + 1,
                    len(self._backends),
                    type(backend).__name__,
                    exc,
                )

        # Should be unreachable (mock always succeeds)
        logger.error("AutoBackend: all %d backends exhausted", len(self._backends))
        raise RuntimeError("All backends exhausted")


# ---------------------------------------------------------------------------
# Riddle generator
# ---------------------------------------------------------------------------
class RiddleGenerator:
    """Orchestrate prompt building, LLM inference, and response validation.

    The LLM draws on its own training knowledge of literary and mythological
    themes — no static corpus is needed.

    Parameters
    ----------
    backend : LLMBackend
        The concrete LLM implementation to use.
    max_retries : int, default 2
        How many times to retry on JSON parse / validation errors.
    """

    def __init__(
        self,
        backend: LLMBackend,
        max_retries: int = MAX_RETRIES,
    ) -> None:
        self._backend = backend
        self._max_retries = max_retries

    def _build_prompt(self, plaintext: str, theme: str) -> str:
        """Construct the full prompt for the LLM.

        No static corpus excerpt is injected — the LLM draws on its
        training knowledge to produce authentic thematically-grounded
        riddles.

        Parameters
        ----------
        plaintext : str
            The secret message to encode.
        theme : str
            Theme key (e.g. ``"greek_myth"``).

        Returns
        -------
        str
            The formatted prompt.
        """
        from alien_obfuscator.config import THEME_LABELS

        theme_label = THEME_LABELS.get(theme, theme)
        prompt = SYSTEM_PROMPT_TEMPLATE.format(
            theme_description=theme_label,
            theme_name=theme_label,
            plaintext=plaintext,
            num_options=NUM_OPTIONS,
            num_distractors=NUM_OPTIONS - 1,
        )
        prompt += STRICT_JSON_PROMPT
        return prompt

    def _parse_response(self, raw: str) -> dict[str, Any]:
        """Clean and parse the raw LLM output into a validated dict.

        Strips markdown fences (```json ... ```) if present. If the full text
        is not valid JSON, attempts to extract the first JSON object ``{...}``
        from within the text (handles some models that wrap JSON in
        chain-of-thought).

        Parameters
        ----------
        raw : str
            Raw text from the LLM.

        Returns
        -------
        dict[str, Any]
            Validated riddle dict.

        Raises
        ------
        ValueError
            If the text cannot be parsed or validated.
        """
        text = raw.strip()
        if not text:
            logger.error("LLM returned empty response")
            raise ValueError("LLM returned empty response")
        if text.startswith("```"):
            first_newline = text.find("\n")
            if first_newline >= 0:
                text = text[first_newline + 1 :]
            else:
                text = text[3:]
            if text.endswith("```"):
                text = text[:-3]
            text = text.strip()

        # Try full text first
        if text.startswith("{"):
            try:
                data = json.loads(text)
                return _validate_riddle_json(data)
            except (json.JSONDecodeError, ValueError):
                pass

        # Fallback: scan for the first valid JSON object in the text
        pos = 0
        while True:
            start = text.find("{", pos)
            if start < 0:
                break
            depth = 0
            matched = False
            for end in range(start, len(text)):
                if text[end] == "{":
                    depth += 1
                elif text[end] == "}":
                    depth -= 1
                    if depth == 0:
                        matched = True
                        candidate = text[start : end + 1]
                        try:
                            data = json.loads(candidate)
                            logger.info("Extracted JSON from text (len=%d)", len(candidate))
                            return _validate_riddle_json(data)
                        except (json.JSONDecodeError, ValueError):
                            pass
                        break
            if not matched:
                pos = start + 1
            else:
                pos = end + 1

        logger.warning(
            "Failed to parse LLM output as JSON: Raw text (len=%d): %.400s",
            len(raw),
            text[:400],
        )
        raise ValueError("Could not extract valid JSON from LLM response.")

    def generate(self, plaintext: str, theme: str) -> dict[str, Any]:
        """Generate a riddle + MCQ options for the given plaintext and theme.

        Retries up to ``max_retries`` times if the LLM returns malformed
        JSON. On success, the ``options`` list is shuffled and
        ``correct_index`` is updated accordingly.

        Parameters
        ----------
        plaintext : str
            The secret message to encode.
        theme : str
            Theme key or ``"surprise"``.

        Returns
        -------
        dict[str, Any]
            A validated riddle dict with shuffled options.

        Raises
        ------
        ValueError
            If the plaintext is empty or exceeds the length limit, or if the
            LLM output cannot be parsed as valid JSON.
        RuntimeError
            If all retries are exhausted without producing valid JSON.
        """
        if not plaintext or not plaintext.strip():
            raise ValueError("Plaintext must not be empty.")
        if len(plaintext) > MAX_PLAINTEXT_LENGTH:
            raise ValueError(f"Plaintext exceeds {MAX_PLAINTEXT_LENGTH} characters.")

        last_error: Exception | None = None

        for _attempt in range(self._max_retries + 1):
            try:
                prompt = self._build_prompt(plaintext, theme)
                raw = self._backend.generate(prompt)
                data = self._parse_response(raw)
                break
            except (json.JSONDecodeError, ValueError, RuntimeError) as exc:
                logger.warning("Attempt %d/%d failed: %s", _attempt + 1, self._max_retries + 1, exc)
                last_error = exc
                continue
        else:
            logger.error("All %d attempts exhausted", self._max_retries + 1)
            raise RuntimeError(
                f"Failed to generate valid riddle after {self._max_retries + 1} attempts."
            ) from last_error

        # Shuffle options so the correct answer moves to a random position
        options = data["options"]
        correct_answer = options[data["correct_index"]]
        random.shuffle(options)
        data["correct_index"] = options.index(correct_answer)
        data["theme"] = theme
        return data