"""Lenient JSON validation for the NEMOCITY mind (one-shot BUILD + FIX).

NEMOCITY has no GBNF grammars: every backend generates free text and the
parsers here recover the JSON. The old `grammar` parameter on the Backend
protocol is repurposed as a plain string TAG telling the backend which shape
is wanted (temperature + validation), never an actual grammar:

    BUILD_TAG      "build"        BUILD permit JSON
    FIX_TAG        "fix"          FIX choice JSON
    MODERATION_TAG "moderation"   moderation verdict JSON
    None                          free text
    "<tag>!strict"                retry pass: temperature 0

Parsing is the same strictness for every backend: mock output is trusted but
parsed identically so the pipeline never special-cases a backend. The engine
does ALL semantic clamping (kind synonyms, floors ranges, placement) — these
parsers only guarantee structure.
"""

from __future__ import annotations

import json
import re

BUILD_TAG = "build"
FIX_TAG = "fix"
MODERATION_TAG = "moderation"
STRICT_SUFFIX = "!strict"

# Name kept from godseed so copied call sites (moderation_judge) still work;
# it is a tag now, not a grammar.
MODERATION_GRAMMAR = MODERATION_TAG

NAME_MAX_CHARS = 24
NEAR_MAX_CHARS = 60
BLURB_MAX_CHARS = 140
DIAGNOSIS_MAX_CHARS = 160
FIX_BLURB_MAX_CHARS = 120
DECLINE_MAX_CHARS = 160
MAX_BUILDINGS = 3


def is_moderation_grammar(grammar: str | None) -> bool:
    return bool(grammar) and (
        str(grammar).split("!", 1)[0] == MODERATION_TAG or '"allowed"' in str(grammar)
    )


def base_tag(grammar: str | None) -> str | None:
    """Tag without the !strict retry suffix."""
    return None if grammar is None else str(grammar).split("!", 1)[0]


def is_strict(grammar: str | None) -> bool:
    return bool(grammar) and str(grammar).endswith(STRICT_SUFFIX)


_THINK_RE = re.compile(r"<think>.*?</think>", re.S)
_FENCE_RE = re.compile(r"```[a-zA-Z]*")


def strip_noise(raw: str) -> str:
    """Drop <think> blocks and markdown code fences before brace-scanning."""
    text = _THINK_RE.sub("", str(raw or ""))
    text = text.replace("<think>", "").replace("</think>", "")
    return _FENCE_RE.sub("", text)


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

    Brace-scans with string/escape awareness so JSON embedded in stray prose
    (a misbehaving unconstrained backend) is still recoverable.
    """
    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 _one_line(text: str, limit: int) -> str:
    return " ".join(str(text or "").split())[:limit]


def _as_int(value) -> int | None:
    if isinstance(value, bool):
        return None
    if isinstance(value, int):
        return value
    if isinstance(value, float):
        return int(round(value))
    if isinstance(value, str):
        try:
            return int(round(float(value.strip())))
        except ValueError:
            return None
    return None


def _coerce_building(item) -> dict | None:
    """One BUILD entry, normalized to a stable 5-key shape; None if unusable.

    kind must be a non-empty string (the engine maps synonyms / unknowns).
    Unknown keys are dropped; types are coerced; name/near are whitespace-
    collapsed and clipped.
    """
    if not isinstance(item, dict):
        return None
    kind = item.get("kind")
    if not isinstance(kind, str) or not kind.strip():
        return None
    name = item.get("name")
    name = _one_line(name, NAME_MAX_CHARS) if isinstance(name, str) and name.strip() else None
    near = item.get("near")
    near = _one_line(near, NEAR_MAX_CHARS) if isinstance(near, str) and near.strip() else None
    return {
        "kind": _one_line(kind.lower(), 40),
        "name": name,
        "near": near,
        "floors": _as_int(item.get("floors")),
        "hue": _as_int(item.get("hue")),
    }


_BUILDING_FRAGMENT_RE = re.compile(r'\{[^{}]*"kind"\s*:\s*"')


def parse_build(raw: str) -> tuple[dict | None, str | None]:
    """Parse a one-shot BUILD reply. Returns (plan, None) or (None, error).

    Plan shape (always fully normalized):
        {"intent": "build"|"decline", "blurb": str,
         "buildings": [entry x0..3], "decline_reason": str|None}

    Lenient by design: code fences and <think> blocks are stripped, the outer
    object is balanced-brace extracted, then salvaged key by key. When the
    outer JSON is broken, individual {"kind": ...} objects are scraped from
    the raw text so a truncated array still yields buildings.
    """
    if not raw or not str(raw).strip():
        return None, "empty output"
    text = strip_noise(str(raw))

    obj = None
    candidate = extract_json_object(text)
    if candidate is not None:
        try:
            obj = json.loads(candidate)
        except (json.JSONDecodeError, ValueError):
            obj = None

    intent = "build"
    blurb = ""
    decline_reason = None
    buildings: list[dict] = []
    if isinstance(obj, dict):
        if obj.get("intent") in ("build", "decline"):
            intent = obj["intent"]
        b = obj.get("blurb")
        if isinstance(b, str):
            blurb = _one_line(b, BLURB_MAX_CHARS)
        dr = obj.get("decline_reason")
        if isinstance(dr, str) and dr.strip():
            decline_reason = _one_line(dr, DECLINE_MAX_CHARS)
        arr = obj.get("buildings")
        if isinstance(arr, list):
            for item in arr:
                entry = _coerce_building(item)
                if entry is not None:
                    buildings.append(entry)

    # Salvage: broken outer object / broken array — scrape balanced
    # {"kind": ...} fragments straight from the text.
    if not buildings:
        for m in _BUILDING_FRAGMENT_RE.finditer(text):
            frag = extract_json_object(text[m.start():])
            if not frag:
                continue
            try:
                item = json.loads(frag)
            except (json.JSONDecodeError, ValueError):
                continue
            entry = _coerce_building(item)
            if entry is not None:
                buildings.append(entry)

    buildings = buildings[:MAX_BUILDINGS]
    if intent == "decline" and not buildings:
        return {
            "intent": "decline",
            "blurb": blurb,
            "buildings": [],
            "decline_reason": decline_reason or "the petition could not be read as a building",
        }, None
    if not buildings:
        return None, "no usable BUILD JSON"
    return {"intent": "build", "blurb": blurb, "buildings": buildings, "decline_reason": None}, None


_FIX_CHOICE_RE = re.compile(r'"choice"\s*:\s*"([^"]{1,12})"')
_FIX_DIAGNOSIS_RE = re.compile(r'"diagnosis"\s*:\s*"([^"]{1,300})"')
_FIX_BLURB_RE = re.compile(r'"blurb"\s*:\s*"([^"]{1,200})"')


def parse_fix(raw: str) -> tuple[dict | None, str | None]:
    """Parse a one-shot FIX reply. Returns (obj, None) or (None, error).

    Obj shape: {"diagnosis": str, "choice": str|None, "blurb": str}.
    choice is NOT validated against the candidate list here — the planner
    does that and falls back to the best predicted candidate.
    """
    if not raw or not str(raw).strip():
        return None, "empty output"
    text = strip_noise(str(raw))

    diagnosis = choice = blurb = None
    candidate = extract_json_object(text)
    if candidate is not None:
        try:
            obj = json.loads(candidate)
        except (json.JSONDecodeError, ValueError):
            obj = None
        if isinstance(obj, dict):
            d = obj.get("diagnosis")
            if isinstance(d, str) and d.strip():
                diagnosis = _one_line(d, DIAGNOSIS_MAX_CHARS)
            c = obj.get("choice")
            if isinstance(c, (str, int)) and not isinstance(c, bool) and str(c).strip():
                choice = str(c).strip()[:12]
            b = obj.get("blurb")
            if isinstance(b, str) and b.strip():
                blurb = _one_line(b, FIX_BLURB_MAX_CHARS)

    if diagnosis is None and choice is None:
        # Key-by-key salvage from a broken outer object.
        m = _FIX_CHOICE_RE.search(text)
        choice = m.group(1).strip() if m else None
        m = _FIX_DIAGNOSIS_RE.search(text)
        diagnosis = _one_line(m.group(1), DIAGNOSIS_MAX_CHARS) if m else None
        if blurb is None:
            m = _FIX_BLURB_RE.search(text)
            blurb = _one_line(m.group(1), FIX_BLURB_MAX_CHARS) if m else None

    if diagnosis is None and choice is None:
        return None, "no usable FIX JSON"
    return {"diagnosis": diagnosis or "", "choice": choice, "blurb": blurb or ""}, None


def parse_moderation(raw: str) -> tuple[dict | None, str | None]:
    """Parse a moderation verdict. Returns (obj, None) or (None, error).

    `allowed` must be a real JSON boolean — anything fuzzier than that is a
    parse error, and parse errors mean DENY upstream (default-deny).
    """
    if not raw or not raw.strip():
        return None, "empty output"
    candidate = extract_json_object(strip_noise(raw))
    if candidate is None:
        return None, "no JSON object found"
    try:
        obj = json.loads(candidate)
    except (json.JSONDecodeError, ValueError) as exc:
        return None, f"invalid JSON: {exc.msg if hasattr(exc, 'msg') else exc}"
    if not isinstance(obj, dict):
        return None, "top level is not a JSON object"
    allowed = obj.get("allowed")
    if not isinstance(allowed, bool):
        return None, '"allowed" is not a boolean'
    category = obj.get("category")
    if category is not None and not isinstance(category, str):
        return None, '"category" is not a string'
    return {"allowed": allowed, "category": (category or "").strip()[:40]}, None