"""Per-field results -> one schema-valid JSON response (SPEC.md §6).

The model runtime produces raw outcomes (head probabilities, a span, generated text); the
helpers here turn them into :class:`FieldResult` values with fixed, deterministic rules
(first-maximum tie-breaks, the jev confidence formula, the parsers in :mod:`deg.parsers`),
and :func:`assemble` builds ``{model, value, meta, usage, plan}`` and validates ``value``
against the input schema.

A result that cannot be represented (a span out of range, generated text breaking its
bounds, a required non-nullable field with nothing extracted) raises :class:`FillError`:
that is a runtime bug, never something to paper over in the response.
"""

from __future__ import annotations

import math
import re
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from decimal import Decimal
from typing import Any

from jsonschema import FormatChecker
from jsonschema import validators as js_validators

from deg.parsers import FORMAT_CHECKS, DateOrder, parse_date, parse_date_time, parse_number
from deg.schema import ArrayNode, FieldPlan, FillPlan, LeafNode, Node, ObjectNode

MODEL_ID = "deg-qwen3.5-0.8b-0.1"


class FillError(ValueError):
    """A per-field result or assembled value violates the plan/schema (HTTP 500)."""

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


class NoValueError(FillError):
    """A required non-nullable (``forbid``) extract field has no valid value in the
    context. Outside anchored array items this is the response (HTTP 500, SPEC.md §4);
    inside one it means the item does not exist and the runtime drops it (§12.6)."""


@dataclass(frozen=True)
class FieldResult:
    """Value for one concrete leaf. ``omit`` drops the key (optional field, nothing found)."""

    value: Any
    meta: dict[str, Any] = field(default_factory=dict)
    omit: bool = False


def concrete_pointer(parts: Sequence[str | int]) -> str:
    """RFC 6901 pointer for a concrete location (array indices are integers)."""
    out = []
    for part in parts:
        out.append(str(part) if isinstance(part, int) else
                   part.replace("~", "~0").replace("/", "~1"))
    return "/" + "/".join(out) if out else ""


def confidence(probabilities: Sequence[float]) -> float:
    """jev formula: (count * peak - 1) / (count - 1), clamped to [0, 1]."""
    count = len(probabilities)
    if count < 2:
        raise ValueError("confidence needs at least two outcomes")
    peak = max(probabilities)
    return min(1.0, max(0.0, (count * peak - 1.0) / (count - 1.0)))


def _normalized(probabilities: Sequence[float], count: int, ptr: str) -> list[float]:
    values = [float(p) for p in probabilities]
    if len(values) != count or not all(math.isfinite(p) and p >= 0.0 for p in values):
        raise FillError([{"loc": ptr, "msg": f"expected {count} finite non-negative "
                                             "probabilities"}])
    total = sum(values)
    if total <= 0.0:
        raise FillError([{"loc": ptr, "msg": "probabilities must have positive mass"}])
    return [p / total for p in values]


def _argmax(values: Sequence[float]) -> int:
    return max(range(len(values)), key=values.__getitem__)  # first maximum on ties


def _expect_mode(fld: FieldPlan, *modes: str) -> None:
    if fld.mode not in modes:
        raise FillError([{"loc": fld.pointer, "msg": f"field mode is {fld.mode}, not {modes}"}])


# ----------------------------------------------------------------------------- decide


def decide_choice(fld: FieldPlan, probabilities: Sequence[float]) -> FieldResult:
    """One probability per option, in schema order."""
    _expect_mode(fld, "decide.choice")
    assert fld.options is not None
    probs = _normalized(probabilities, len(fld.options), fld.pointer)
    best = _argmax(probs)
    return FieldResult(
        fld.options[best],
        {"mode": fld.mode, "probability": probs[best], "confidence": confidence(probs),
         "probabilities": probs},
    )


def decide_noul(fld: FieldPlan, p_true: float) -> FieldResult:
    """``true`` iff p_true >= 0.5."""
    _expect_mode(fld, "decide.noul")
    p = float(p_true)
    if not (math.isfinite(p) and 0.0 <= p <= 1.0):
        raise FillError([{"loc": fld.pointer, "msg": "noul probability must be in [0, 1]"}])
    return FieldResult(p >= 0.5, {"mode": fld.mode, "probability": p})


def decide_score(fld: FieldPlan, probabilities: Sequence[float]) -> FieldResult:
    """One probability per integer level (+ a trailing null outcome when ``null_outcome``).

    The value is the most probable level; ``meta["expected"]`` is the expected level. A
    winning null outcome yields ``null``.
    """
    _expect_mode(fld, "decide.score")
    assert fld.levels is not None
    count = len(fld.levels) + (1 if fld.null_outcome else 0)
    probs = _normalized(probabilities, count, fld.pointer)
    best = _argmax(probs)
    meta: dict[str, Any] = {"mode": fld.mode, "probability": probs[best],
                            "confidence": confidence(probs), "probabilities": probs}
    level_probs = probs[: len(fld.levels)]
    mass = sum(level_probs)
    if mass > 0.0:
        meta["expected"] = sum(lv * p for lv, p in zip(fld.levels, level_probs)) / mass
    if fld.null_outcome and best == len(fld.levels):
        return FieldResult(None, meta)
    return FieldResult(fld.levels[best], meta)


# ----------------------------------------------------------------------------- extract


def _string_violation(fld: FieldPlan, value: str) -> str | None:
    c = fld.constraints
    if len(value) < c.get("minLength", 0):
        return "minLength"
    if "maxLength" in c and len(value) > c["maxLength"]:
        return "maxLength"
    if "pattern" in c and not re.search(c["pattern"], value):
        return "pattern"
    if fld.format is not None and not FORMAT_CHECKS[fld.format](value):
        return "format"
    return None


def _number_violation(fld: FieldPlan, value: float) -> str | None:
    c = fld.constraints
    if "minimum" in c and value < c["minimum"]:
        return "minimum"
    if "maximum" in c and value > c["maximum"]:
        return "maximum"
    if "exclusiveMinimum" in c and value <= c["exclusiveMinimum"]:
        return "exclusiveMinimum"
    if "exclusiveMaximum" in c and value >= c["exclusiveMaximum"]:
        return "exclusiveMaximum"
    if "multipleOf" in c:
        q = Decimal(repr(value)) / Decimal(repr(c["multipleOf"]))
        if q != q.to_integral_value():
            return "multipleOf"
    return None


def _null(fld: FieldPlan, meta: dict[str, Any], reason: str) -> FieldResult:
    meta = {**meta, "null_reason": reason}
    if fld.on_null == "forbid":
        raise NoValueError([{"loc": fld.pointer,
                             "msg": f"required non-nullable field has no value ({reason})"}])
    return FieldResult(None, meta, omit=fld.on_null == "omit")


def extract(
    fld: FieldPlan,
    context_text: str,
    span: tuple[int, int] | None,
    span_confidence: float,
    *,
    date_order: DateOrder | None = None,
    default_offset: str = "Z",
) -> FieldResult:
    """Copy ``context_text[start:end]`` (code-point offsets) and parse it per mode.

    ``span=None`` means "not in the context". A span that fails its parser or constraints
    is treated the same way, with ``null_reason`` saying why.
    """
    _expect_mode(fld, "extract.span", "extract.number", "extract.date")
    c = float(span_confidence)
    if not (math.isfinite(c) and 0.0 <= c <= 1.0):
        raise FillError([{"loc": fld.pointer, "msg": "confidence must be in [0, 1]"}])
    meta: dict[str, Any] = {"mode": fld.mode, "confidence": c, "span": None}
    if span is None:
        return _null(fld, meta, "not_found")
    start, end = span
    if not (isinstance(start, int) and isinstance(end, int)
            and 0 <= start < end <= len(context_text)):
        raise FillError([{"loc": fld.pointer, "msg": f"span {span!r} out of range"}])
    meta["span"] = [start, end]
    text = context_text[start:end]
    if fld.mode == "extract.span":
        violation = _string_violation(fld, text)
        return _null(fld, meta, violation) if violation else FieldResult(text, meta)
    if fld.mode == "extract.number":
        number = parse_number(text, integer=fld.json_type == "integer")
        if number is None:
            return _null(fld, meta, "unparsable")
        violation = _number_violation(fld, number)
        return _null(fld, meta, violation) if violation else FieldResult(number, meta)
    if fld.format == "date":
        day = parse_date(text, date_order=date_order)
        if day is None:
            return _null(fld, meta, "unparsable")
        violation = _string_violation(fld, day)
        return _null(fld, meta, violation) if violation else FieldResult(day, meta)
    parsed = parse_date_time(text, default_offset=default_offset)
    if parsed is None:
        return _null(fld, meta, "unparsable")
    stamp, assumed = parsed
    meta["assumed_offset"] = assumed
    violation = _string_violation(fld, stamp)
    return _null(fld, meta, violation) if violation else FieldResult(stamp, meta)


# ----------------------------------------------------------------------------- generate


def generate(fld: FieldPlan, text: str, generated_tokens: int) -> FieldResult:
    """Generated text must already satisfy minLength/maxLength/pattern (bounded decoding)."""
    _expect_mode(fld, "generate")
    if not isinstance(generated_tokens, int) or generated_tokens < 0:
        raise FillError([{"loc": fld.pointer, "msg": "generated_tokens must be an int >= 0"}])
    violation = _string_violation(fld, text)
    if violation:
        raise FillError([{"loc": fld.pointer,
                          "msg": f"generated text violates {violation}; decoding must be "
                                 "bounded by the field's constraints"}])
    return FieldResult(text, {"mode": fld.mode, "generated_tokens": generated_tokens})


# ----------------------------------------------------------------------------- assemble


def format_checker() -> FormatChecker:
    """Formats apply to strings only (JSON Schema); other instance types pass."""
    checker = FormatChecker(formats=())
    for name, check in FORMAT_CHECKS.items():
        checker.checks(name)(lambda value, check=check: not isinstance(value, str) or check(value))
    return checker


def validate_value(plan: FillPlan, value: Any) -> None:
    """Validate an assembled value against the plan's input schema (+ format checks)."""
    cls = js_validators.validator_for(plan.schema, default=js_validators.Draft202012Validator)
    validator = cls(plan.schema, format_checker=format_checker())
    errors = sorted(validator.iter_errors(value), key=lambda e: list(map(str, e.absolute_path)))
    if errors:
        raise FillError([{"loc": concrete_pointer(list(e.absolute_path)), "msg": e.message}
                         for e in errors])


def assemble(
    plan: FillPlan,
    results: Mapping[str, FieldResult],
    array_lengths: Mapping[str, int],
    *,
    prompt_tokens: int,
    model: str = MODEL_ID,
    dropped_items: Mapping[str, int] | None = None,
) -> dict[str, Any]:
    """Build the response body.

    ``results`` and ``array_lengths`` are keyed by concrete RFC 6901 pointer (e.g.
    ``/line_items/0/amount``). Every leaf the plan implies must have exactly one result;
    unused keys are an error. ``dropped_items`` (array pointer -> n > 0) adds
    ``"dropped": n`` to that array's meta: array items removed because a required
    value was not in the context (SPEC.md §12.6).
    """
    if not isinstance(prompt_tokens, int) or prompt_tokens < 0:
        raise FillError([{"loc": "usage", "msg": "prompt_tokens must be an int >= 0"}])
    meta: dict[str, dict[str, Any]] = {}
    used_results: set[str] = set()
    used_arrays: set[str] = set()
    errors: list[dict[str, Any]] = []

    def walk(node: Node, parts: list[str | int]) -> Any:
        ptr = concrete_pointer(parts)
        if isinstance(node, LeafNode):
            result = results.get(ptr)
            if result is None:
                errors.append({"loc": ptr, "msg": "missing field result"})
                return None
            used_results.add(ptr)
            if result.meta.get("mode") != node.field.mode:
                errors.append({"loc": ptr, "msg": f"meta.mode must be {node.field.mode}"})
            meta[ptr] = result.meta
            return result
        if isinstance(node, ArrayNode):
            length = array_lengths.get(ptr)
            if not isinstance(length, int) or isinstance(length, bool) or length < 0:
                errors.append({"loc": ptr, "msg": "missing or invalid array length"})
                return []
            used_arrays.add(ptr)
            meta[ptr] = {"mode": "array", "items": length}
            if dropped_items and dropped_items.get(ptr):
                meta[ptr]["dropped"] = dropped_items[ptr]
            items = []
            for index in range(length):
                item = walk(node.item, [*parts, index])
                items.append(item.value if isinstance(item, FieldResult) else item)
            return items
        assert isinstance(node, ObjectNode)
        out: dict[str, Any] = {}
        for name, child in node.properties:
            got = walk(child, [*parts, name])
            if isinstance(got, FieldResult):
                if not got.omit:
                    out[name] = got.value
            else:
                out[name] = got
        return out

    value = walk(plan.root, [])
    for ptr in sorted(set(results) - used_results):
        errors.append({"loc": ptr, "msg": "result for a field the plan does not contain"})
    for ptr in sorted(set(array_lengths) - used_arrays):
        errors.append({"loc": ptr, "msg": "length for an array the plan does not contain"})
    if errors:
        raise FillError(errors)
    validate_value(plan, value)
    generated = sum(int(m.get("generated_tokens", 0)) for m in meta.values())
    return {
        "model": model,
        "value": value,
        "meta": meta,
        "usage": {"prompt_tokens": prompt_tokens, "generated_tokens": generated},
        "plan": plan.fingerprint,
    }