Blame for jue_flash/encoding/render.py

"""Text rendering of state, instructions, and option criteria for Jue-Flash.

Includes:
1. Explicit 1-based cardinality indexing ``[i/N]`` on multi-element JSON lists so bidirectional encoders
   can directly read list lengths and hop counts without recurrent counting.
2. Automatic taxonomy enrichment from the embedded 92 KB MITRE CWE / ATT&CK glossary
   (``glossary.json.gz``) when candidate options specify bare ``CWE-xxx`` / ``Txxxx`` / ``Gxxxx`` IDs
   with ``None`` or empty descriptions.
"""

from __future__ import annotations

import gzip
import json
from functools import lru_cache
from pathlib import Path
from typing import Union

from jue_flash.api.schema import ChoiceQuestion, NoulQuestion, Question, ScoreQuestion

JSONContent = Union[str, dict, list, int, float, bool, None]

PRIMITIVE_TO_ID = {"choice": 0, "score": 1, "noul": 2}
ID_TO_PRIMITIVE = {0: "choice", 1: "score", 2: "noul"}

_GLOSSARY_PATH = Path(__file__).with_name("glossary.json.gz")


@lru_cache(maxsize=1)
def load_taxonomy_glossary() -> dict[str, str]:
    """Load the embedded MITRE CWE / ATT&CK glossary dictionary."""
    if not _GLOSSARY_PATH.exists():
        return {}
    with gzip.open(_GLOSSARY_PATH, "rt", encoding="utf-8") as f:
        return json.load(f)


def render(v: JSONContent, indent: int = 0) -> str:
    """Flatten str | object | array into structured text. Field names and array indices are preserved."""
    pad = "  " * indent
    if v is None:
        return ""
    if isinstance(v, (str, int, float, bool)):
        return str(v)
    if isinstance(v, list):
        n = len(v)
        if n == 0:
            return f"{pad}[empty_list len=0]"
        if n == 1:
            return f"{pad}- {render(v[0], indent + 1).lstrip()}"
        return "\n".join(f"{pad}[{i}/{n}] {render(x, indent + 1).lstrip()}" for i, x in enumerate(v, 1))
    return "\n".join(
        f"{pad}{k}:\n{render(x, indent + 1)}" if isinstance(x, (dict, list)) else f"{pad}{k}: {render(x)}"
        for k, x in v.items()
    )


def option_text(name: str, desc: JSONContent) -> str:
    """Render a single candidate option/level/noul pole into text, enriching bare CWE/ATT&CK keys."""
    rendered = render(desc).strip() if desc is not None else ""
    if not rendered:
        glossary = load_taxonomy_glossary()
        norm_key = name.strip().upper()
        if norm_key in glossary:
            return f"{name}: {glossary[norm_key]}"
        return name
    return f"{name}: {rendered}"


def render_question_parts(question: Question) -> tuple[str, list[tuple[str, str]]]:
    """Return ``(instruction_text, [(option_key, option_rendered_text), ...])`` in canonical readout order."""
    inst = render(question.instructions).strip() if question.instructions is not None else ""
    if not inst:
        inst = f"Evaluate {question.type} decision."

    if isinstance(question, ChoiceQuestion):
        opts = [(k, option_text(k, desc)) for k, desc in question.criteria.items()]
        return inst, opts
    if isinstance(question, ScoreQuestion):
        opts = [(str(lvl), option_text(str(lvl), desc)) for lvl, desc in enumerate(question.criteria)]
        return inst, opts
    if isinstance(question, NoulQuestion):
        false_desc = question.criteria.false_ if question.criteria is not None else None
        true_desc = question.criteria.true_ if question.criteria is not None else None
        opts = [
            ("false", option_text("false", false_desc if false_desc is not None else "No / False / Absent")),
            ("true", option_text("true", true_desc if true_desc is not None else "Yes / True / Present")),
        ]
        return inst, opts
    raise TypeError(f"Unsupported question type: {type(question).__name__}")