mind/schema.py

"""Awakening JSON schema, defensive parsing, and the single-retry repair prompt.

The séance (ARCHITECTURE.md §3) asks MiniCPM-V for ONE structured JSON object
per awakening. The model has no native JSON mode, so parsing is defensive by
contract: strip code-fence lines, brace-balance-scan for the first complete
``{...}`` (string- and escape-aware — never a greedy regex), validate hard with
pydantic, then run the writing agent's line validators. Any failure raises
:class:`AwakeningParseError` with a message terse enough to feed straight back
to the model via :func:`build_repair_prompt` — backends get exactly ONE repair
retry before giving up poetically (:class:`PoeticError`).

The "her" pattern, applied: the model owns WHAT it saw and how the soul talks;
this module owns whether the reply is structurally true to the contract.
"""

from __future__ import annotations

import inspect
import json
import logging
import re
from typing import Literal, Mapping, Optional

from pydantic import BaseModel, Field, ValidationError, model_validator

logger = logging.getLogger("pareidolia.mind.schema")

# ---------------------------------------------------------------------------
# Line validators (writing agent's mind/prompts.py owns the comedy bar).
#
# Contract (per prompts.py): LINE_VALIDATORS is a mapping (or iterable) of
# callables with signature
#
#     validator(line: str, kind: str) -> str | None
#
# where ``kind`` is "grudge" or "mutter". Return None/"" when the line
# passes; return a short error string when it fails (the string rides back to
# the model in the repair prompt). One-arg validators ``validator(line)`` are
# also supported. Import-guarded so the mind package tests and the mock
# backend run standalone even without prompts.py.
# ---------------------------------------------------------------------------
try:  # pragma: no cover - exercised only once prompts.py exists
    from .prompts import LINE_VALIDATORS  # type: ignore[attr-defined]
except Exception:  # noqa: BLE001 - any import problem means "no validators yet"
    LINE_VALIDATORS: tuple = ()  # type: ignore[no-redef]


class PoeticError(RuntimeError):
    """A failure whose message is safe to show a visitor verbatim.

    Raised when the spirits genuinely decline (double parse failure, model
    misbehavior). ``str(err)`` is finished user-facing copy — never a stack
    trace, never an internal detail. Internal context travels on
    ``err.__cause__`` for the logs.
    """

    def __init__(self, reason: str):
        super().__init__(reason)
        self.reason = reason


class AwakeningParseError(ValueError):
    """The model's reply did not satisfy the §3 contract.

    The message is written FOR the model: short, specific, actionable — it is
    interpolated into :func:`build_repair_prompt` for the single retry.
    """


# ---------------------------------------------------------------------------
# Pydantic models — the §3 awakening shape
# ---------------------------------------------------------------------------

FeatureRole = Literal["eye_left", "eye_right", "mouth"]


class Gate(BaseModel):
    """Moderation booleans the VLM must return; the server enforces them.

    Layered-moderation rule (§0): nothing reaches CV snap, TTS, or the public
    wall when :meth:`refusal` returns copy.
    """

    contains_human_face: bool
    nsfw: bool
    recognizable_object: bool

    def refusal(self) -> Optional[str]:
        """Poetic refusal copy when the gate closes, else None.

        Checked in severity order; human faces win because that copy is the
        funniest true thing we can say about a person.
        """
        if self.contains_human_face:
            return "It is already awake."
        if self.nsfw:
            return "The spirits decline."
        if not self.recognizable_object:
            return "The mist found nothing to hold onto. Step closer, or try another angle."
        return None


class Feature(BaseModel):
    """One EXISTING visual element of the photo, named as seen.

    Two grounding modes:

    - **marker** (preferred): ``marker`` is the id of a numbered badge that
      cv/propose.py painted onto a real detected feature of the image.
      ``cv.propose.resolve_markers`` later replaces ``cx``/``cy``/``size``
      with that anchor's exact values (``anchor_kind="marker"``). Models pick
      marked regions far better than they emit numbers — this kills
      coordinate hallucination at the root.
    - **coarse** (fallback): normalized [0,1] coordinates relative to the
      image (cx right, cy down); ``size`` is a normalized diameter. These are
      the VLM's coarse guesses — cv/snap.py later moves them to the strongest
      nearby anchor and records ``snap_delta``. The mist animation forgives
      ±15%, so coarse is fine; out-of-range is not.

    A feature must carry a ``marker`` or the full coarse triple (cx, cy,
    size) — enforced below so an unplaceable feature never reaches the
    overlay.
    """

    name: str = Field(min_length=1, max_length=80)
    role: FeatureRole
    marker: Optional[int] = Field(default=None, ge=1)
    cx: Optional[float] = Field(default=None, ge=0.0, le=1.0)
    cy: Optional[float] = Field(default=None, ge=0.0, le=1.0)
    size: Optional[float] = Field(default=None, gt=0.0, le=1.0)

    @model_validator(mode="after")
    def _grounded(self) -> "Feature":
        if self.marker is None:
            missing = [k for k in ("cx", "cy", "size") if getattr(self, k) is None]
            if missing:
                raise ValueError(
                    'feature needs either "marker" (a numbered badge id) or '
                    f"all of cx, cy, size (missing: {', '.join(missing)})"
                )
        return self


class Persona(BaseModel):
    """Who the object turns out to have been all along.

    ``archetype`` and ``voice`` come from the banks in WRITING.md; they stay
    plain strings here because the archetype bank belongs to prompts.py and
    voice.py degrades gracefully on an unknown voice id (taste lives there,
    not in the schema).
    """

    archetype: str = Field(min_length=1, max_length=40)
    voice: str = Field(min_length=1, max_length=40)
    mood: str = Field(min_length=1, max_length=60)


class Lines(BaseModel):
    """The two lines the object gets to say.

    Structural caps only — the comedy rules (≤22 words, visible specifics,
    no "I am a…") are enforced by LINE_VALIDATORS from prompts.py so the
    writing agent owns the bar in exactly one place.
    """

    grudge: str = Field(min_length=1, max_length=280)
    mutter: str = Field(min_length=1, max_length=160)


class AwakeningResult(BaseModel):
    """The full §3 awakening record — one VLM generation, validated.

    Field notes:
    - ``object``/``material``/``condition``/``setting`` ground every line in
      THIS photo (condition's canonical vocabulary is rusted|chipped|pristine|
      abandoned|worn|dusty|broken|loved but stays free-text — the archetype
      bank, not the schema, decides what a condition means).
    - ``candidate_features`` must carry exactly one eye_left and one
      eye_right (the first blink is the hero moment) and at most one mouth.
    - ``critique`` is the model's mandatory self-check, written BEFORE the
      finals in generation order; requiring it keeps the
      critique-then-final contract honest.
    """

    gate: Gate
    object: str = Field(min_length=1, max_length=80)
    material: str = Field(min_length=1, max_length=80)
    condition: str = Field(min_length=1, max_length=40)
    setting: str = Field(min_length=1, max_length=120)
    candidate_features: list[Feature] = Field(min_length=2, max_length=6)
    critique: str = Field(min_length=1, max_length=400)
    persona: Persona
    lines: Lines

    @model_validator(mode="after")
    def _roles_complete(self) -> "AwakeningResult":
        roles = [f.role for f in self.candidate_features]
        for eye in ("eye_left", "eye_right"):
            if roles.count(eye) != 1:
                raise ValueError(
                    f'candidate_features needs exactly one "{eye}" '
                    f"(got {roles.count(eye)})"
                )
        if roles.count("mouth") > 1:
            raise ValueError('candidate_features allows at most one "mouth"')
        return self


# ---------------------------------------------------------------------------
# Defensive extraction
# ---------------------------------------------------------------------------

# Full-line code fences only (```json / ```). JSON strings cannot contain a
# raw newline, so a whole-line match can never sit inside the payload.
_FENCE_LINE = re.compile(r"^\s*```[a-zA-Z0-9_-]*\s*$", re.MULTILINE)


def _strip_code_fences(text: str) -> str:
    return _FENCE_LINE.sub("", text)


def extract_json_object(raw: str) -> Optional[str]:
    """Return the first balanced top-level ``{...}`` in ``raw``, or None.

    Brace-scans with string/escape awareness so JSON wrapped in prose, fences,
    or trailing chatter is still recoverable. Truncated output (depth never
    returns to 0) yields None — which is exactly what we want: a truncated
    generation must trigger the repair retry, not a partial parse.
    """
    start = raw.find("{")
    if start < 0:
        return None
    depth = 0
    in_string = False
    escaped = False
    for i in range(start, len(raw)):
        ch = raw[i]
        if in_string:
            if escaped:
                escaped = False
            elif ch == "\\":
                escaped = True
            elif ch == '"':
                in_string = False
            continue
        if ch == '"':
            in_string = True
        elif ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                return raw[start : i + 1]
    return None


def _format_validation_error(exc: ValidationError) -> str:
    """Squash a pydantic error into one short line the model can act on."""
    parts = []
    for err in exc.errors()[:3]:
        loc = ".".join(str(p) for p in err["loc"]) or "(root)"
        parts.append(f"{loc}: {err['msg']}")
    more = len(exc.errors()) - 3
    if more > 0:
        parts.append(f"(+{more} more)")
    return "; ".join(parts)


def _iter_line_validators():
    validators = LINE_VALIDATORS
    if isinstance(validators, Mapping):
        return list(validators.values())
    return list(validators)


def _call_validator(validator, line: str, kind: str) -> Optional[str]:
    """Invoke a line validator with (line, kind) when its signature allows, else (line)."""
    try:
        params = [
            p
            for p in inspect.signature(validator).parameters.values()
            if p.kind
            in (
                inspect.Parameter.POSITIONAL_ONLY,
                inspect.Parameter.POSITIONAL_OR_KEYWORD,
                inspect.Parameter.VAR_POSITIONAL,
            )
        ]
        wants_two = len(params) >= 2 or any(
            p.kind is inspect.Parameter.VAR_POSITIONAL for p in params
        )
    except (TypeError, ValueError):  # builtins / odd callables
        wants_two = False
    return validator(line, kind) if wants_two else validator(line)


# Minimum lateral (cx) separation between the two eyes. The VLM sometimes
# returns eye_left/eye_right nearly (or exactly) on top of each other — a
# degenerate pair that renders as one smeared eye and kills the first-blink
# hero moment. Rejecting here (not in the pydantic model) keeps the message
# on the repair-retry path while the image is still in context.
MIN_EYE_SEPARATION_CX = 0.05

EYE_SEPARATION_ERROR = (
    "eye features must be two distinct, laterally separated elements; "
    "pick two different visible features"
)

EYE_MARKER_ERROR = (
    "eye_left and eye_right use the same marker — eyes need two DIFFERENT "
    "markers; pick another numbered badge for one of them"
)


def _eye_separation_error(result: AwakeningResult) -> Optional[str]:
    """The degenerate-eye-pair error for this result, or None when it passes.

    Marker-grounded eyes (either eye carries a ``marker``) are judged on
    marker identity instead of raw cx: distinct anchors are guaranteed at
    least 0.06 apart by cv/propose.py's dedupe, while a shared marker is one
    smeared eye. Coarse pairs keep the MIN_EYE_SEPARATION_CX lateral check.
    Returned (not raised) so parse_awakening can soften the failure on the
    repair retry.
    """
    by_role = {f.role: f for f in result.candidate_features}
    left, right = by_role.get("eye_left"), by_role.get("eye_right")
    if left is None or right is None:  # pragma: no cover - model enforces this
        return None
    if left.marker is not None or right.marker is not None:
        if left.marker is not None and left.marker == right.marker:
            return EYE_MARKER_ERROR
        return None  # distinct anchors are deduped >=0.06 apart upstream
    if left.cx is None or right.cx is None:  # pragma: no cover - _grounded forbids
        return None
    if abs(left.cx - right.cx) < MIN_EYE_SEPARATION_CX:
        return EYE_SEPARATION_ERROR
    return None


def _unknown_marker_error(
    result: AwakeningResult, allowed_markers: Optional[set[int]]
) -> Optional[str]:
    """The unknown-marker-id error for this result, or None when it passes.

    Only enforced when ``allowed_markers`` is given (the marks-grounded flow);
    legacy coarse-coordinate parses pass None and skip this entirely. The
    message names the valid ids so the repair retry can actually fix it.
    """
    if allowed_markers is None:
        return None
    valid = {int(m) for m in allowed_markers}
    unknown = sorted(
        {f.marker for f in result.candidate_features if f.marker is not None and f.marker not in valid}
    )
    if not unknown:
        return None
    bad = ", ".join(str(m) for m in unknown)
    if valid:
        ids = ", ".join(str(m) for m in sorted(valid))
        return (
            f"unknown marker id(s): {bad} — the valid marker ids are {ids}; "
            'choose from the numbered badges in the image, or omit "marker" '
            "and give coarse cx, cy, size instead"
        )
    return (
        f"unknown marker id(s): {bad} — this image has no numbered markers; "
        'omit "marker" and give coarse cx, cy, size instead'
    )


def _run_line_validators(result: AwakeningResult) -> None:
    for kind, line in (("grudge", result.lines.grudge), ("mutter", result.lines.mutter)):
        for validator in _iter_line_validators():
            error = _call_validator(validator, line, kind)
            if error:
                raise AwakeningParseError(f"lines.{kind} rejected: {error}")


def parse_awakening(
    text: str,
    allowed_markers: Optional[set[int]] = None,
    *,
    retry: bool = False,
) -> AwakeningResult:
    """Parse one raw VLM reply into a validated :class:`AwakeningResult`.

    Pipeline: strip fence lines -> brace-balance extract the first {...} ->
    json.loads -> pydantic validate -> unknown-marker check (only when
    ``allowed_markers`` is given) -> eye-separation check ->
    LINE_VALIDATORS (prompts.py). Raises
    :class:`AwakeningParseError` whose message is ready for
    :func:`build_repair_prompt`; never raises raw pydantic/json errors.

    ``allowed_markers``: the set of badge ids cv/propose.py painted on the
    image (marks-grounded flow). Every feature ``marker`` must be in it; the
    rejection message names the valid ids so the repair retry can fix it.
    None (legacy coarse flow) skips marker validation entirely.

    ``retry``: pass True on the SECOND parse — the repair retry. The single
    behavior change: an eye-separation-only failure is ACCEPTED with a logged
    warning instead of raising, because web/overlay.js's render-side spread
    clamp already displays pairs <0.05 cx apart safely, and a flat-but-
    servable record beats a PoeticError. First-parse behavior is unchanged
    (the failure still drives the repair retry while the image is in
    context); every other failure class still raises on the retry.
    """
    if not text or not text.strip():
        raise AwakeningParseError("empty reply — output the JSON object")
    candidate = extract_json_object(_strip_code_fences(text))
    if candidate is None:
        raise AwakeningParseError(
            "no complete JSON object found (reply may be truncated or prose-only)"
        )
    try:
        payload = json.loads(candidate)
    except (json.JSONDecodeError, ValueError) as exc:
        msg = exc.msg if isinstance(exc, json.JSONDecodeError) else str(exc)
        raise AwakeningParseError(f"invalid JSON: {msg}") from exc
    if not isinstance(payload, dict):
        raise AwakeningParseError("top level must be a JSON object")
    try:
        result = AwakeningResult.model_validate(payload)
    except ValidationError as exc:
        raise AwakeningParseError(_format_validation_error(exc)) from exc
    marker_error = _unknown_marker_error(result, allowed_markers)
    if marker_error is not None:
        raise AwakeningParseError(marker_error)
    eye_error = _eye_separation_error(result)
    if eye_error is not None:
        if retry:
            logger.warning(
                "eye separation soft-accepted on the repair retry "
                "(render-side spread clamp handles it): %s",
                eye_error,
            )
        else:
            raise AwakeningParseError(eye_error)
    _run_line_validators(result)
    return result


# ---------------------------------------------------------------------------
# Face geometry gate — post-snap plausibility for the live coarse+snap flow
# ---------------------------------------------------------------------------

# June 12 live audit: the VLM places a plausible face ~57% raw, and the
# failure modes are GEOMETRIC and detectable on the FINAL (post-snap) points:
# eye pairs cramped onto one feature cluster (stand mixer: both eyes tight on
# the bowl rim), eyes drifting off level in busy texture (bonsai foliage),
# and borderline edge/glint picks (toaster chrome, cone edge). Each rule
# below yields a short human sentence that rides VERBATIM into the geometry
# repair prompt, so thresholds live here next to the words that explain them.

EYE_LEVEL_MAX_DCY = 0.08  # eyes further apart vertically than this read tilted
EYE_CRAMPED_MIN_SEP = 0.10  # lateral floor — stricter than the 0.05 parse gate
EYE_CRAMPED_SIZE_FACTOR = 1.2  # big eyes need proportionally more separation
MOUTH_BELOW_FACTOR = 0.5  # mouth must clear BOTH eyes by this x mean eye size
MOUTH_SPAN_SLACK = 0.15  # lateral slack beyond the eye span for the mouth
CENTRAL_REGION = (0.08, 0.92)  # cx/cy outside this is an off-object edge pick


def _geom_value(feature: Mapping, key: str) -> Optional[float]:
    """The numeric ``feature[key]`` as float, or None when absent/non-numeric."""
    value = feature.get(key)
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        return None
    return float(value)


def face_geometry_violations(features) -> list[str]:
    """Geometry violations of the FINAL (post-snap) features, as sentences.

    PURE function over plain feature dicts (``role``, ``cx``/``cy``/``size``
    normalized to [0,1]) — call it AFTER cv snap, on what will actually
    render. Each returned string is a short human sentence ready for
    :func:`build_geometry_repair_prompt`. Any check whose role (the mouth is
    optional per the schema) or coordinates are missing is skipped, so a
    clean face — and a face we cannot judge — both return ``[]``.
    """
    by_role: dict = {}
    for feature in features or []:
        if not isinstance(feature, Mapping):
            continue
        role = feature.get("role")
        if isinstance(role, str) and role not in by_role:
            by_role[role] = feature

    left, right, mouth = (
        by_role.get(role) for role in ("eye_left", "eye_right", "mouth")
    )
    violations: list[str] = []

    eye_sizes = [
        size
        for eye in (left, right)
        if eye is not None
        for size in (_geom_value(eye, "size"),)
        if size is not None
    ]
    mean_eye_size = sum(eye_sizes) / len(eye_sizes) if eye_sizes else 0.0

    cx_l = _geom_value(left, "cx") if left is not None else None
    cx_r = _geom_value(right, "cx") if right is not None else None
    cy_l = _geom_value(left, "cy") if left is not None else None
    cy_r = _geom_value(right, "cy") if right is not None else None

    if cy_l is not None and cy_r is not None and abs(cy_l - cy_r) > EYE_LEVEL_MAX_DCY:
        violations.append(
            f"the eyes are not level: eye_left sits at height {cy_l:.2f} but "
            f"eye_right at {cy_r:.2f}; pick two features at roughly the same height"
        )

    if cx_l is not None and cx_r is not None:
        min_sep = max(EYE_CRAMPED_MIN_SEP, EYE_CRAMPED_SIZE_FACTOR * mean_eye_size)
        sep = abs(cx_l - cx_r)
        if sep < min_sep:
            violations.append(
                f"the eyes are cramped onto one spot (only {sep:.2f} apart; "
                f"they need at least {min_sep:.2f} of the image width between them)"
            )

    if mouth is not None:
        m_cx = _geom_value(mouth, "cx")
        m_cy = _geom_value(mouth, "cy")
        if (
            m_cy is not None
            and cy_l is not None
            and cy_r is not None
            and m_cy - max(cy_l, cy_r) < MOUTH_BELOW_FACTOR * mean_eye_size
        ):
            violations.append(
                "the mouth is not clearly below both eyes; pick a mouth "
                "feature that sits lower on the object than both eyes"
            )
        if m_cx is not None and cx_l is not None and cx_r is not None:
            lo = min(cx_l, cx_r) - MOUTH_SPAN_SLACK
            hi = max(cx_l, cx_r) + MOUTH_SPAN_SLACK
            if not (lo <= m_cx <= hi):
                violations.append(
                    "the mouth sits off to the side, outside the span of the "
                    "eyes; pick a mouth feature between the eyes"
                )

    region_lo, region_hi = CENTRAL_REGION
    for feature in features or []:
        if not isinstance(feature, Mapping):
            continue
        cx = _geom_value(feature, "cx")
        cy = _geom_value(feature, "cy")
        if any(
            value is not None and not (region_lo <= value <= region_hi)
            for value in (cx, cy)
        ):
            role = feature.get("role")
            name = feature.get("name")
            label = str(role) if isinstance(role, str) and role else "feature"
            if isinstance(name, str) and name:
                label = f"{label} ({name})"
            violations.append(
                f"the {label} sits at the very edge of the frame, likely off "
                "the object; pick a feature nearer the object's center"
            )

    return violations


def build_geometry_repair_prompt(violations: list[str]) -> str:
    """Build the geometry complaint for the multi-turn repair retry.

    Same philosophy as :func:`build_repair_prompt`: name the precise
    problems, demand bare JSON. Sent as a follow-up user turn with the image
    still in context (the backend replays the accepted exchange), listing the
    violations verbatim. The backend keeps the FIRST accepted parse's
    persona/lines no matter what comes back — asking the model not to touch
    them just keeps the reply focused on features.
    """
    listed = "\n".join(f"- {violation}" for violation in violations)
    return (
        "The face you found does not sit right on the object:\n"
        f"{listed}\n"
        "Look at the image again and pick different visible features for the "
        "failing roles; keep everything else (object, persona, lines) the "
        "same. Reply with ONLY the corrected JSON object — same schema, no "
        "code fences, no prose before or after the braces."
    )


# ---------------------------------------------------------------------------
# Repair prompt — the single retry
# ---------------------------------------------------------------------------

_REPAIR_ECHO_MAX = 1500  # bound the prefill; the tail is where truncation bites


def build_repair_prompt(error: str, original_text: str) -> str:
    """Build the one-shot repair prompt fed back with the same image.

    Mirrors godseed's validate-retry approach: name the precise problem, echo
    the (clipped) previous output, demand bare JSON. One retry only — a model
    that fails twice gets a PoeticError upstream, not a third chance.
    """
    clipped = (original_text or "").strip()
    if len(clipped) > _REPAIR_ECHO_MAX:
        clipped = clipped[:_REPAIR_ECHO_MAX] + " …"
    return (
        "Your previous reply could not be accepted.\n"
        f"Problem: {error}\n"
        "Your previous reply was:\n"
        "---\n"
        f"{clipped}\n"
        "---\n"
        "Look at the image again and reply with ONLY the corrected JSON object "
        "— same schema, no code fences, no prose before or after the braces. "
        "Fix the problem named above; keep everything that was already good."
    )