"""JSON Schema validator and compiler: schema -> deterministic fill plan.

The mode of every leaf is decided purely by its JSON type and keywords (see SPEC.md §3);
there is no router and no model choice. The only author-controlled switch is
``x-generate: true``, allowed on ``string`` and on ``array`` of ``string``.

Compilation is a pure function of the schema: the same schema always yields the same plan,
and a leaf's :attr:`FieldPlan.fingerprint` depends only on that leaf (its key path, its own
keywords and annotations), never on its siblings or on property order elsewhere.

Paths are tuples of property names, with :data:`ITEM` (``None``) standing for "each item of
the enclosing array". :attr:`FieldPlan.pointer` renders them as a JSON-Pointer-like string
with ``*`` for :data:`ITEM` (a literal ``*`` key is escaped as ``~2``).

Invalid or unsupported schemas raise :class:`SchemaError` carrying every problem found as a
list of ``{"loc", "msg"}`` (HTTP 422 in the planned server).
"""

from __future__ import annotations

import hashlib
import json
import math
import re
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field, replace
from typing import Any, Literal

from jsonschema import exceptions as js_exceptions
from jsonschema import validators as js_validators

Mode = Literal[
    "decide.choice",
    "decide.score",
    "decide.noul",
    "extract.span",
    "extract.number",
    "extract.date",
    "generate",
]
DECIDE_MODES: frozenset[str] = frozenset({"decide.choice", "decide.score", "decide.noul"})
EXTRACT_MODES: frozenset[str] = frozenset({"extract.span", "extract.number", "extract.date"})

GENERATE_KEY = "x-generate"
MAX_TOKENS_KEY = "x-max-tokens"  # per-field decode budget (x-generate fields only)
LEVELS_KEY = "x-levels"  # score level labels {"<level>": "label"} (decide.score only)
EXTENSIONS = frozenset({GENERATE_KEY, MAX_TOKENS_KEY, LEVELS_KEY})
ITEM = None  # path element: "each item of the enclosing array"

MIN_OPTIONS = 2
MAX_OPTIONS = 255  # same limit as jev choice questions
MAX_SCORE_LEVELS = 11  # integers with <= 11 admissible values (0..10, 1..5) are decided;
# every other number (any ``number``, or a wider/open integer range) is extracted

DATE_FORMATS: frozenset[str] = frozenset({"date", "date-time"})
SPAN_FORMATS: frozenset[str] = frozenset({"email", "uri", "uuid"})

ANNOTATIONS: frozenset[str] = frozenset(
    {
        "title",
        "description",
        "default",
        "examples",
        "$comment",
        "$schema",
        "$id",
        "$anchor",
        "$defs",
        "definitions",
        "readOnly",
        "writeOnly",
        "deprecated",
    }
)
STRING_KEYWORDS = frozenset({"minLength", "maxLength", "pattern", "format"})
NUMBER_KEYWORDS = frozenset(
    {"minimum", "maximum", "exclusiveMinimum", "exclusiveMaximum", "multipleOf"}
)
OBJECT_KEYWORDS = frozenset({"properties", "required", "additionalProperties"})
ARRAY_KEYWORDS = frozenset({"items", "minItems", "maxItems", "uniqueItems"})
TYPE_KEYWORDS: dict[str, frozenset[str]] = {
    "string": STRING_KEYWORDS,
    "integer": NUMBER_KEYWORDS,
    "number": NUMBER_KEYWORDS,
    "boolean": frozenset(),
    "object": OBJECT_KEYWORDS,
    "array": ARRAY_KEYWORDS,
}
COMMON_KEYWORDS = frozenset({"type", "enum", "$ref", "anyOf", "oneOf"}) | EXTENSIONS

# Explicitly unsupported keywords and why. Anything else not listed above is "unknown".
UNSUPPORTED: dict[str, str] = {
    "allOf": "schema intersection is not supported; merge the subschemas",
    "not": "negated schemas cannot be filled",
    "if": "conditional schemas are not supported",
    "then": "conditional schemas are not supported",
    "else": "conditional schemas are not supported",
    "const": "a single fixed value has nothing to decide; omit the field",
    "dependentSchemas": "conditional schemas are not supported",
    "dependentRequired": "conditional requirements are not supported",
    "dependencies": "conditional requirements are not supported",
    "patternProperties": "objects must list their properties explicitly",
    "propertyNames": "objects must list their properties explicitly",
    "unevaluatedProperties": "objects must list their properties explicitly",
    "minProperties": "objects must list their properties explicitly",
    "maxProperties": "objects must list their properties explicitly",
    "unevaluatedItems": "arrays must have one homogeneous items schema",
    "prefixItems": "tuple arrays are not supported; use an object",
    "additionalItems": "tuple arrays are not supported; use an object",
    "contains": "arrays must have one homogeneous items schema",
    "minContains": "arrays must have one homogeneous items schema",
    "maxContains": "arrays must have one homogeneous items schema",
    "$dynamicRef": "only local $ref ('#/...') is supported",
    "$dynamicAnchor": "only local $ref ('#/...') is supported",
    "$recursiveRef": "only local $ref ('#/...') is supported",
    "$recursiveAnchor": "only local $ref ('#/...') is supported",
    "contentEncoding": "encoded content is not supported",
    "contentMediaType": "encoded content is not supported",
    "contentSchema": "encoded content is not supported",
}


class SchemaError(ValueError):
    """Invalid or unsupported schema; ``errors`` is a list of ``{loc, msg}``."""

    def __init__(self, errors: list[dict[str, Any]], status: int = 422) -> None:
        super().__init__("; ".join(f"{'/'.join(map(str, e['loc']))}: {e['msg']}" for e in errors))
        self.errors = errors
        self.status = status

    def body(self) -> dict[str, Any]:
        return {"error": {"type": "invalid_schema", "status": self.status, "detail": self.errors}}


Path = tuple[str | None, ...]


def pointer(path: Path) -> str:
    """Render a plan path; ``ITEM`` -> ``*``; ``~``, ``/``, ``*`` escaped as ~0, ~1, ~2."""
    parts = []
    for part in path:
        if part is ITEM:
            parts.append("*")
        else:
            parts.append(part.replace("~", "~0").replace("/", "~1").replace("*", "~2"))
    return "/" + "/".join(parts) if parts else ""


def canonical_json(value: Any) -> str:
    return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False)


def _digest(value: Any) -> str:
    return hashlib.sha256(canonical_json(value).encode()).hexdigest()


@dataclass(frozen=True, eq=True)
class FieldPlan:
    """One leaf to fill. ``constraints`` holds the JSON Schema keywords the value must meet.

    * ``options``: ``decide.choice`` outcomes (JSON values, schema order); with
      ``option_descriptions`` from ``oneOf``/``anyOf`` const branches when present.
    * ``levels``: ``decide.score`` outcomes (the admissible integers, at most 11).
    * ``level_labels``: per-level label text from ``x-levels`` (aligned with ``levels``).
    * ``null_outcome``: decide heads get an extra "null" outcome (nullable score fields).
    * ``max_new_tokens``: a generate field's own decode budget (``x-max-tokens``).
    * ``on_null``: what an extract field does when nothing is found: ``"null"`` (nullable),
      ``"omit"`` (optional, non-nullable) or ``"forbid"`` (required, non-nullable: the
      extractor must return its best span).
    """

    path: Path
    mode: Mode
    json_type: str
    required: bool
    nullable: bool
    title: str | None = None
    description: str | None = None
    options: tuple[Any, ...] | None = None
    option_descriptions: tuple[str | None, ...] | None = None
    levels: tuple[int, ...] | None = None
    level_labels: tuple[str | None, ...] | None = None
    null_outcome: bool = False
    max_new_tokens: int | None = None
    format: str | None = None
    constraints: dict[str, Any] = field(default_factory=dict)

    @property
    def pointer(self) -> str:
        return pointer(self.path)

    @property
    def on_null(self) -> Literal["null", "omit", "forbid"]:
        if self.nullable:
            return "null"
        return "forbid" if self.required else "omit"

    @property
    def phase(self) -> Literal["parallel", "generate"]:
        return "generate" if self.mode == "generate" else "parallel"

    def to_dict(self) -> dict[str, Any]:
        out: dict[str, Any] = {
            "pointer": self.pointer,
            "path": list(self.path),
            "mode": self.mode,
            "type": self.json_type,
            "required": self.required,
            "nullable": self.nullable,
            "on_null": self.on_null,
        }
        for key in ("title", "description", "format"):
            if getattr(self, key) is not None:
                out[key] = getattr(self, key)
        if self.options is not None:
            out["options"] = list(self.options)
        if self.option_descriptions is not None:
            out["option_descriptions"] = list(self.option_descriptions)
        if self.levels is not None:
            out["levels"] = list(self.levels)
        if self.level_labels is not None:
            out["level_labels"] = list(self.level_labels)
        if self.max_new_tokens is not None:
            out["max_new_tokens"] = self.max_new_tokens
        if self.null_outcome:
            out["null_outcome"] = True
        if self.constraints:
            out["constraints"] = dict(self.constraints)
        return out

    @property
    def fingerprint(self) -> str:
        """sha256 of the leaf's canonical spec: identical wherever this leaf appears."""
        return _digest(self.to_dict())


@dataclass(frozen=True, eq=True)
class ArrayPlan:
    """An array node. Its length is a runtime count bounded by min/max items."""

    path: Path
    min_items: int
    max_items: int | None
    unique_items: bool
    required: bool

    @property
    def pointer(self) -> str:
        return pointer(self.path)

    def to_dict(self) -> dict[str, Any]:
        return {
            "pointer": self.pointer,
            "path": list(self.path),
            "min_items": self.min_items,
            "max_items": self.max_items,
            "unique_items": self.unique_items,
            "required": self.required,
        }


@dataclass(frozen=True)
class LeafNode:
    field: FieldPlan


@dataclass(frozen=True)
class ArrayNode:
    array: ArrayPlan
    item: Node


@dataclass(frozen=True)
class ObjectNode:
    path: Path
    properties: tuple[tuple[str, Node], ...]


Node = LeafNode | ArrayNode | ObjectNode


@dataclass(frozen=True)
class FillPlan:
    """Compiled schema. ``fields`` are leaves in schema order (depth-first)."""

    root: ObjectNode
    fields: tuple[FieldPlan, ...]
    arrays: tuple[ArrayPlan, ...]
    schema: Mapping[str, Any]

    @property
    def parallel_fields(self) -> tuple[FieldPlan, ...]:
        """decide.* and extract.*: independent of each other and of order."""
        return tuple(f for f in self.fields if f.phase == "parallel")

    @property
    def generate_fields(self) -> tuple[FieldPlan, ...]:
        """generate: run after the parallel phase, in this (schema) order."""
        return tuple(f for f in self.fields if f.phase == "generate")

    def field(self, ptr: str) -> FieldPlan:
        for f in self.fields:
            if f.pointer == ptr:
                return f
        raise KeyError(ptr)

    def to_dict(self) -> dict[str, Any]:
        return {
            "fields": [f.to_dict() for f in self.fields],
            "arrays": [a.to_dict() for a in self.arrays],
            "fingerprint": self.fingerprint,
        }

    @property
    def fingerprint(self) -> str:
        """Order-insensitive over the parallel phase, order-sensitive over generate fields.

        Reordering decide/extract properties keeps the fingerprint; reordering generate
        fields changes it, because each generated field conditions on the ones before it.
        """
        return _digest(
            {
                "parallel": sorted(f.fingerprint for f in self.parallel_fields),
                "generate": [f.fingerprint for f in self.generate_fields],
                "arrays": sorted(_digest(a.to_dict()) for a in self.arrays),
            }
        )


# ----------------------------------------------------------------------------- compiler


def _is_number(value: Any) -> bool:
    return isinstance(value, (int, float)) and not isinstance(value, bool)


def _type_matches(value: Any, json_type: str) -> bool:
    if json_type == "null":
        return value is None
    if json_type == "boolean":
        return isinstance(value, bool)
    if json_type == "integer":
        return _is_number(value) and float(value).is_integer()
    if json_type == "number":
        return _is_number(value) and math.isfinite(value)
    if json_type == "string":
        return isinstance(value, str)
    return False


class _Compiler:
    def __init__(self, root: Mapping[str, Any]) -> None:
        self.root = root
        self.errors: list[dict[str, Any]] = []
        self.fields: list[FieldPlan] = []
        self.arrays: list[ArrayPlan] = []

    def error(self, loc: Sequence[Any], msg: str) -> None:
        self.errors.append({"loc": list(loc), "msg": msg})

    # -- $ref / anyOf / oneOf normalisation -------------------------------------------

    def resolve(
        self, schema: Any, loc: list[Any], refs: tuple[str, ...]
    ) -> tuple[dict[str, Any] | None, tuple[str, ...]]:
        """Follow local ``$ref`` (annotation siblings override the target's)."""
        while isinstance(schema, dict) and "$ref" in schema:
            ref = schema["$ref"]
            if not isinstance(ref, str) or not (ref == "#" or ref.startswith("#/")):
                self.error([*loc, "$ref"], f"only local $ref ('#/...') is supported, got {ref!r}")
                return None, refs
            if ref in refs:
                self.error([*loc, "$ref"], f"recursive $ref {ref!r} is not supported")
                return None, refs
            extra = set(schema) - {"$ref"} - ANNOTATIONS - EXTENSIONS
            if extra:
                self.error(loc, f"$ref siblings must be annotations, got {sorted(extra)}")
                return None, refs
            target: Any = self.root
            for raw in ref[2:].split("/") if ref != "#" else []:
                part = raw.replace("~1", "/").replace("~0", "~")
                if isinstance(target, dict) and part in target:
                    target = target[part]
                elif isinstance(target, list) and part.isdigit() and int(part) < len(target):
                    target = target[int(part)]
                else:
                    self.error([*loc, "$ref"], f"unresolvable $ref {ref!r}")
                    return None, refs
            if not isinstance(target, dict):
                self.error([*loc, "$ref"], f"$ref {ref!r} must point at a schema object")
                return None, refs
            siblings = {k: v for k, v in schema.items() if k != "$ref"}
            schema = {**target, **siblings}
            refs = (*refs, ref)
        if not isinstance(schema, dict):
            self.error(loc, "schema must be an object (boolean schemas are not supported)")
            return None, refs
        return schema, refs

    def flatten_union(
        self, schema: dict[str, Any], loc: list[Any], refs: tuple[str, ...]
    ) -> tuple[dict[str, Any] | None, tuple[str, ...], list[str | None] | None]:
        """anyOf/oneOf allowed only as (a) nullable ``[X, {type: null}]`` or (b) an enum
        written as const/enum branches (branch descriptions become option descriptions)."""
        key = "anyOf" if "anyOf" in schema else "oneOf" if "oneOf" in schema else None
        if key is None:
            return schema, refs, None
        if "anyOf" in schema and "oneOf" in schema:
            self.error(loc, "anyOf and oneOf together are not supported")
            return None, refs, None
        extra = set(schema) - {key} - ANNOTATIONS - {"type"}
        if extra:
            self.error(loc, f"{key} siblings must be annotations, got {sorted(extra)}")
            return None, refs, None
        raw_branches = schema[key]
        if not isinstance(raw_branches, list) or not raw_branches:
            self.error([*loc, key], f"{key} must be a non-empty array")
            return None, refs, None
        branches = []
        for index, raw in enumerate(raw_branches):
            branch, _ = self.resolve(raw, [*loc, key, index], refs)
            if branch is None:
                return None, refs, None
            branches.append(branch)
        annotations = {k: v for k, v in schema.items() if k in ANNOTATIONS}

        def is_null(b: dict[str, Any]) -> bool:
            return b.get("type") == "null" and not (set(b) - ANNOTATIONS - {"type"})

        nulls = [b for b in branches if is_null(b)]
        others = [b for b in branches if not is_null(b)]
        if len(nulls) == 1 and len(others) == 1:
            other = others[0]
            if any(k in other for k in ("anyOf", "oneOf")):
                self.error([*loc, key], f"nested {key} is not supported")
                return None, refs, None
            merged = {**other, **annotations}
            base = other.get("type")
            if isinstance(base, str):
                merged["type"] = [base, "null"]
            elif "enum" in other and isinstance(other["enum"], list):
                merged["enum"] = [*other["enum"], *([] if None in other["enum"] else [None])]
            else:
                self.error([*loc, key], "nullable branch must declare a single type or enum")
                return None, refs, None
            return merged, refs, None
        options: list[Any] = []
        descriptions: list[str | None] = []
        for index, branch in enumerate(branches):
            values = None
            if "const" in branch and not (set(branch) - ANNOTATIONS - {"const", "type"}):
                values = [branch["const"]]
            elif "enum" in branch and not (set(branch) - ANNOTATIONS - {"enum", "type"}):
                values = branch["enum"] if isinstance(branch["enum"], list) else None
            elif is_null(branch):
                values = [None]
            if values is None:
                self.error(
                    [*loc, key, index],
                    f"{key} is supported only for nullable [X, {{type: null}}] or for "
                    "branches that are all const/enum (an enum with descriptions)",
                )
                return None, refs, None
            desc = branch.get("description", branch.get("title"))
            for value in values:
                options.append(value)
                descriptions.append(desc if isinstance(desc, str) else None)
        merged = {**annotations, "enum": options}
        if "type" in schema:
            merged["type"] = schema["type"]
        if all(d is None for d in descriptions):
            return merged, refs, None
        return merged, refs, descriptions

    # -- nodes ----------------------------------------------------------------------

    def node(
        self,
        raw: Any,
        path: Path,
        loc: list[Any],
        *,
        required: bool,
        generate: bool,
        refs: tuple[str, ...] = (),
        max_tokens: int | None = None,
    ) -> Node | None:
        schema, refs = self.resolve(raw, loc, refs)
        if schema is None:
            return None
        schema, refs, option_descriptions = self.flatten_union(schema, loc, refs)
        if schema is None:
            return None
        before = len(self.errors)
        for key in schema:
            if key in UNSUPPORTED:
                self.error([*loc, key], f"unsupported keyword {key!r}: {UNSUPPORTED[key]}")
            elif key.startswith("x-") and key not in EXTENSIONS:
                self.error([*loc, key], f"unknown extension {key!r} (known: "
                                        f"{', '.join(sorted(EXTENSIONS))})")
        if len(self.errors) > before:
            return None

        flag = schema.get(GENERATE_KEY, False)
        if not isinstance(flag, bool):
            self.error([*loc, GENERATE_KEY], "x-generate must be a boolean")
            return None
        generate = generate or flag
        own_budget = schema.get(MAX_TOKENS_KEY)
        if own_budget is not None:
            if not generate:
                self.error([*loc, MAX_TOKENS_KEY], "x-max-tokens applies only to x-generate "
                                                   "fields")
                return None
            if not isinstance(own_budget, int) or isinstance(own_budget, bool) \
                    or own_budget < 1:
                self.error([*loc, MAX_TOKENS_KEY], "x-max-tokens must be an integer >= 1")
                return None
            max_tokens = own_budget

        json_type, nullable = self.parse_type(schema, loc)
        if json_type is None:
            return None
        allowed = COMMON_KEYWORDS | ANNOTATIONS | TYPE_KEYWORDS.get(json_type, frozenset())
        unknown = sorted(set(schema) - allowed)
        for key in unknown:
            known = any(key in kws for kws in TYPE_KEYWORDS.values())
            msg = (
                f"keyword {key!r} does not apply to type {json_type!r}"
                if known
                else f"unknown keyword {key!r}"
            )
            self.error([*loc, key], msg)
        if unknown:
            return None

        if generate and not (json_type == "string" or json_type == "array"):
            self.error(
                [*loc, GENERATE_KEY],
                f"x-generate is allowed only on string or array of string, not {json_type!r}"
                + (" (enum fields are decided)" if "enum" in schema else ""),
            )
            return None

        if json_type in ("object", "array") and nullable:
            self.error([*loc, "type"], "nullable objects/arrays are not supported")
            return None
        if json_type == "object":
            if "enum" in schema:
                self.error([*loc, "enum"], "enum of objects is not supported")
                return None
            return self.object_node(schema, path, loc, refs)
        if json_type == "array":
            if "enum" in schema:
                self.error([*loc, "enum"], "enum of arrays is not supported")
                return None
            return self.array_node(schema, path, loc, required=required, generate=generate,
                                   refs=refs, max_tokens=max_tokens)
        leaf = self.leaf(schema, json_type, nullable, path, loc, required=required,
                         generate=generate, option_descriptions=option_descriptions)
        if leaf is None:
            return None
        if max_tokens is not None:
            leaf = replace(leaf, max_new_tokens=max_tokens)
        if LEVELS_KEY in schema:
            leaf = self.level_labels(leaf, schema[LEVELS_KEY], [*loc, LEVELS_KEY])
            if leaf is None:
                return None
        self.fields.append(leaf)
        return LeafNode(leaf)

    def parse_type(self, schema: dict[str, Any], loc: list[Any]) -> tuple[str | None, bool]:
        raw = schema.get("type")
        if raw is None:
            enum = schema.get("enum")
            if isinstance(enum, list) and enum:
                kinds = {_enum_kind(v) for v in enum if v is not None}
                if len(kinds) == 1 and None not in kinds:
                    return kinds.pop(), None in enum
                if not kinds:
                    self.error([*loc, "enum"], "enum must contain a non-null value")
                    return None, False
                return "enum", None in enum  # mixed scalar enum
            self.error([*loc, "type"], "type is required (or an enum)")
            return None, False
        types = [raw] if isinstance(raw, str) else raw
        if not isinstance(types, list) or not all(isinstance(t, str) for t in types):
            self.error([*loc, "type"], "type must be a string or array of strings")
            return None, False
        nullable = "null" in types
        rest = [t for t in types if t != "null"]
        if len(rest) != 1 or len(types) != len(set(types)):
            self.error(
                [*loc, "type"],
                f"union types are not supported (only [X, 'null']), got {types!r}",
            )
            return None, False
        if rest[0] not in TYPE_KEYWORDS:
            self.error([*loc, "type"], f"unknown type {rest[0]!r}")
            return None, False
        return rest[0], nullable

    def object_node(
        self, schema: dict[str, Any], path: Path, loc: list[Any], refs: tuple[str, ...]
    ) -> ObjectNode | None:
        props = schema.get("properties", {})
        if not isinstance(props, dict):
            self.error([*loc, "properties"], "properties must be an object")
            return None
        required = schema.get("required", [])
        if not isinstance(required, list) or not all(isinstance(r, str) for r in required):
            self.error([*loc, "required"], "required must be an array of strings")
            return None
        if len(set(required)) != len(required):
            self.error([*loc, "required"], "required must not repeat names")
        for name in required:
            if name not in props:
                self.error([*loc, "required"], f"required property {name!r} is not in properties")
        extra = schema.get("additionalProperties", False)
        if not isinstance(extra, bool):
            self.error(
                [*loc, "additionalProperties"],
                "additionalProperties must be true or false (a schema cannot be filled: the "
                "property names are unknown); deg never emits undeclared properties",
            )
        children: list[tuple[str, Node]] = []
        for name, sub in props.items():
            child = self.node(
                sub, (*path, name), [*loc, "properties", name],
                required=name in required, generate=False, refs=refs,
            )
            if child is not None:
                children.append((name, child))
        return ObjectNode(path, tuple(children))

    def array_node(
        self,
        schema: dict[str, Any],
        path: Path,
        loc: list[Any],
        *,
        required: bool,
        generate: bool,
        refs: tuple[str, ...],
        max_tokens: int | None = None,
    ) -> ArrayNode | None:
        if "items" not in schema:
            self.error([*loc, "items"], "arrays need an items schema")
            return None
        if isinstance(schema["items"], list):
            self.error([*loc, "items"], "tuple arrays (items: [...]) are not supported")
            return None
        min_items = schema.get("minItems", 0)
        max_items = schema.get("maxItems")
        ok = True
        for key, value in (("minItems", min_items), ("maxItems", max_items)):
            if value is not None and (
                not isinstance(value, int) or isinstance(value, bool) or value < 0
            ):
                self.error([*loc, key], f"{key} must be a non-negative integer")
                ok = False
        if ok and max_items is not None and max_items < min_items:
            self.error([*loc, "maxItems"], "maxItems must be >= minItems")
            ok = False
        unique = schema.get("uniqueItems", False)
        if not isinstance(unique, bool):
            self.error([*loc, "uniqueItems"], "uniqueItems must be a boolean")
            ok = False
        if not ok:
            return None
        item_path = (*path, ITEM)
        if generate:
            item_schema, _ = self.resolve(schema["items"], [*loc, "items"], refs)
            if item_schema is None:
                return None
            item_type = item_schema.get("type")
            if item_type != "string" or "enum" in item_schema or "format" in item_schema:
                self.error(
                    [*loc, GENERATE_KEY],
                    "x-generate is allowed only on string or array of string "
                    "(items must be plain type: string)",
                )
                return None
        array = ArrayPlan(path, min_items, max_items, unique, required)
        self.arrays.append(array)
        item = self.node(
            schema["items"], item_path, [*loc, "items"],
            required=True, generate=generate, refs=refs, max_tokens=max_tokens,
        )
        if item is None:
            return None
        return ArrayNode(array, item)

    def level_labels(self, leaf: FieldPlan, raw: Any, loc: list[Any]) -> FieldPlan | None:
        """``x-levels``: {"<level>": "label"} for a decide.score field. Each label becomes
        that level's option line (``<level>: <label>``), the layout the score head was
        trained on (jev levels); unlabeled levels keep the bare number."""
        if leaf.mode != "decide.score" or leaf.levels is None:
            self.error(loc, "x-levels applies only to integer score fields (bounded integers "
                            f"with <= {MAX_SCORE_LEVELS} values)")
            return None
        if not isinstance(raw, dict) or not raw:
            self.error(loc, "x-levels must be a non-empty object {\"<level>\": \"label\"}")
            return None
        labels: dict[int, str] = {}
        for key, label in raw.items():
            try:
                level = int(key)
            except ValueError:
                level = None
            if level is None or str(level) != key or level not in leaf.levels:
                self.error([*loc, key], f"{key!r} is not one of the levels {list(leaf.levels)}")
                return None
            if not isinstance(label, str) or not label.strip():
                self.error([*loc, key], "a level label must be a non-empty string")
                return None
            labels[level] = label
        return replace(leaf, level_labels=tuple(labels.get(v) for v in leaf.levels))

    def leaf(
        self,
        schema: dict[str, Any],
        json_type: str,
        nullable: bool,
        path: Path,
        loc: list[Any],
        *,
        required: bool,
        generate: bool,
        option_descriptions: list[str | None] | None,
    ) -> FieldPlan | None:
        notes = {}
        for key in ("title", "description"):
            value = schema.get(key)
            if value is not None and not isinstance(value, str):
                self.error([*loc, key], f"{key} must be a string")
                return None
            notes[key] = value
        base = dict(path=path, json_type=json_type, required=required, nullable=nullable,
                    **notes)

        if "enum" in schema:
            if generate:
                self.error([*loc, GENERATE_KEY],
                           "x-generate is allowed only on string or array of string "
                           "(enum fields are decided)")
                return None
            return self.choice(schema, base, loc, json_type, option_descriptions)

        if json_type == "boolean":
            if nullable:
                return FieldPlan(mode="decide.choice", options=(True, False, None), **base)
            return FieldPlan(mode="decide.noul", **base)

        if json_type in ("integer", "number"):
            return self.numeric(schema, base, loc, json_type, nullable)

        # string
        constraints = self.string_constraints(schema, loc)
        if constraints is None:
            return None
        fmt = schema.get("format")
        if fmt is not None:
            if generate:
                self.error([*loc, GENERATE_KEY],
                           "x-generate cannot be combined with format (formatted values are "
                           "extracted from the context)")
                return None
            if fmt in DATE_FORMATS:
                return FieldPlan(mode="extract.date", format=fmt, constraints=constraints, **base)
            if fmt in SPAN_FORMATS:
                return FieldPlan(mode="extract.span", format=fmt, constraints=constraints, **base)
            self.error([*loc, "format"],
                       f"unsupported format {fmt!r}; supported: "
                       f"{sorted(DATE_FORMATS | SPAN_FORMATS)}")
            return None
        mode: Mode = "generate" if generate else "extract.span"
        return FieldPlan(mode=mode, constraints=constraints, **base)

    def choice(
        self,
        schema: dict[str, Any],
        base: dict[str, Any],
        loc: list[Any],
        json_type: str,
        option_descriptions: list[str | None] | None,
    ) -> FieldPlan | None:
        enum = schema["enum"]
        if not isinstance(enum, list):
            self.error([*loc, "enum"], "enum must be an array")
            return None
        seen: list[str] = []
        for index, value in enumerate(enum):
            key = canonical_json(value)
            if key in seen:
                self.error([*loc, "enum", index], f"duplicate enum value {value!r}")
                return None
            seen.append(key)
            if isinstance(value, (dict, list)):
                self.error([*loc, "enum", index], "enum values must be scalars")
                return None
            if value is None:
                if not base["nullable"] and "type" in schema:
                    self.error([*loc, "enum", index], "null in enum needs a nullable type")
                    return None
                continue
            if json_type != "enum" and not _type_matches(value, json_type):
                self.error([*loc, "enum", index],
                           f"enum value {value!r} does not match type {json_type!r}")
                return None
        if not MIN_OPTIONS <= len(enum) <= MAX_OPTIONS:
            self.error([*loc, "enum"],
                       f"enum needs {MIN_OPTIONS}..{MAX_OPTIONS} values "
                       "(a single value has nothing to decide)")
            return None
        extra = sorted(set(schema) & (STRING_KEYWORDS | NUMBER_KEYWORDS))
        if extra:
            self.error([*loc, extra[0]], f"{extra} cannot be combined with enum")
            return None
        nullable = None in enum  # JSON Schema: null must be listed to be valid
        base = {**base, "nullable": nullable}
        return FieldPlan(
            mode="decide.choice",
            options=tuple(enum),
            option_descriptions=tuple(option_descriptions) if option_descriptions else None,
            **base,
        )

    def numeric(
        self,
        schema: dict[str, Any],
        base: dict[str, Any],
        loc: list[Any],
        json_type: str,
        nullable: bool,
    ) -> FieldPlan | None:
        constraints: dict[str, Any] = {}
        for key in NUMBER_KEYWORDS:
            if key in schema:
                value = schema[key]
                if not _is_number(value) or not math.isfinite(value):
                    self.error([*loc, key], f"{key} must be a finite number")
                    return None
                constraints[key] = value
        if "multipleOf" in constraints and constraints["multipleOf"] <= 0:
            self.error([*loc, "multipleOf"], "multipleOf must be > 0")
            return None
        lo_in, hi_in = constraints.get("minimum"), constraints.get("maximum")
        lo_ex, hi_ex = constraints.get("exclusiveMinimum"), constraints.get("exclusiveMaximum")
        if json_type == "integer":
            lows = [math.ceil(lo_in)] if lo_in is not None else []
            lows += [math.floor(lo_ex) + 1] if lo_ex is not None else []
            highs = [math.floor(hi_in)] if hi_in is not None else []
            highs += [math.ceil(hi_ex) - 1] if hi_ex is not None else []
            if lows and highs:
                lo, hi = max(lows), min(highs)
                step = constraints.get("multipleOf")
                # Enumerate only small ranges (a wide range is a quantity: extracted).
                if (hi - lo) / (step or 1) > MAX_SCORE_LEVELS * 2:
                    return FieldPlan(mode="extract.number", constraints=constraints, **base)
                levels = [v for v in range(lo, hi + 1)
                          if step is None or _multiple(v, step)]
                if not levels:
                    self.error(loc, "integer bounds admit no value")
                    return None
                if len(levels) == 1:
                    self.error(loc, "integer bounds admit a single value; nothing to decide")
                    return None
                if len(levels) <= MAX_SCORE_LEVELS:
                    return FieldPlan(mode="decide.score", levels=tuple(levels),
                                     null_outcome=nullable, constraints=constraints, **base)
            return FieldPlan(mode="extract.number", constraints=constraints, **base)
        # Every other number is a quantity read from the text: bounds only validate it.
        lows = [v for v in (lo_in, lo_ex) if v is not None]
        highs = [v for v in (hi_in, hi_ex) if v is not None]
        if lows and highs:
            lo, hi = max(lows), min(highs)
            empty = lo > hi or (lo == hi and (lo_ex is not None and lo_ex >= hi
                                              or hi_ex is not None and hi_ex <= lo))
            if empty:
                self.error(loc, "number bounds admit no value")
                return None
        return FieldPlan(mode="extract.number", constraints=constraints, **base)

    def string_constraints(self, schema: dict[str, Any], loc: list[Any]) -> dict[str, Any] | None:
        constraints: dict[str, Any] = {}
        for key in ("minLength", "maxLength"):
            if key in schema:
                value = schema[key]
                if not isinstance(value, int) or isinstance(value, bool) or value < 0:
                    self.error([*loc, key], f"{key} must be a non-negative integer")
                    return None
                constraints[key] = value
        if constraints.get("maxLength", math.inf) < constraints.get("minLength", 0):
            self.error([*loc, "maxLength"], "maxLength must be >= minLength")
            return None
        if "pattern" in schema:
            if not isinstance(schema["pattern"], str):
                self.error([*loc, "pattern"], "pattern must be a string")
                return None
            try:
                re.compile(schema["pattern"])
            except re.error as exc:
                self.error([*loc, "pattern"], f"invalid pattern: {exc}")
                return None
            constraints["pattern"] = schema["pattern"]
        if "format" in schema and not isinstance(schema["format"], str):
            self.error([*loc, "format"], "format must be a string")
            return None
        return constraints


def _enum_kind(value: Any) -> str:
    if isinstance(value, bool):
        return "boolean"
    if isinstance(value, int):
        return "integer"
    if isinstance(value, float):
        return "number"
    if isinstance(value, str):
        return "string"
    return "enum"


def _multiple(value: float, step: float) -> bool:
    quotient = value / step
    return math.isclose(quotient, round(quotient), rel_tol=0.0, abs_tol=1e-9)


def check_schema(schema: Any) -> None:
    """Raise :class:`SchemaError` unless ``schema`` is a valid JSON Schema document."""
    if not isinstance(schema, dict):
        raise SchemaError([{"loc": ["schema"], "msg": "schema must be a JSON object"}])
    cls = js_validators.validator_for(schema, default=js_validators.Draft202012Validator)
    try:
        cls.check_schema(schema)
    except js_exceptions.SchemaError as exc:
        loc = ["schema", *exc.path]
        raise SchemaError([{"loc": loc, "msg": f"invalid JSON Schema: {exc.message}"}]) from exc


def compile_schema(schema: Any) -> FillPlan:
    """Validate ``schema`` and compile it into a :class:`FillPlan`.

    The root must be an object. Raises :class:`SchemaError` listing every problem found.
    """
    check_schema(schema)
    compiler = _Compiler(schema)
    root = compiler.node(schema, (), ["schema"], required=True, generate=False)
    if compiler.errors:
        raise SchemaError(compiler.errors)
    if not isinstance(root, ObjectNode):
        raise SchemaError([{"loc": ["schema", "type"], "msg": "the root schema must be "
                                                               "type: object"}])
    return FillPlan(root, tuple(compiler.fields), tuple(compiler.arrays), schema)