"""Structured output defense — 4 layers (D-020).

Layer 1: response_format={"type":"json_object"} request (auto-degrades on 400
          inside the provider).
Layer 2: prompt-embedded schema hint ("Respond with ONLY valid JSON...").
Layer 3: defensive parse — strip markdown fences, extract first balanced
          JSON object, pydantic model_validate.
Layer 4: single bounded retry with the validation error fed back.
"""

from typing import TypeVar

from pydantic import BaseModel, ValidationError

from ..llm.base import LLMProvider
from ..llm.types import Message

T = TypeVar("T", bound=BaseModel)


class StructuredOutputError(Exception):
    """Raised when the model output cannot be validated after one retry."""


def extract_json_object(text: str) -> str:
    """Strip fences and return the first balanced {...} block from text."""
    stripped = text.strip()
    if stripped.startswith("```"):
        first_newline = stripped.find("\n")
        if first_newline != -1:
            stripped = stripped[first_newline + 1:]
        if stripped.rstrip().endswith("```"):
            stripped = stripped.rstrip()[:-3]
        stripped = stripped.strip()
    start = stripped.find("{")
    if start == -1:
        raise StructuredOutputError("no JSON object found in model output")
    depth = 0
    in_string = False
    escape = False
    for i, ch in enumerate(stripped[start:], start=start):
        if escape:
            escape = False
            continue
        if ch == "\\":
            escape = True
            continue
        if ch == '"' and not escape:
            in_string = not in_string
            continue
        if in_string:
            continue
        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                return stripped[start:i + 1]
    raise StructuredOutputError("unbalanced JSON object in model output")


def parse_structured(text: str, schema: type[T]) -> T:
    """Layer 3: fence-strip + first-balanced-object + pydantic validation."""
    candidate = extract_json_object(text)
    try:
        return schema.model_validate_json(candidate)
    except ValidationError as exc:
        raise StructuredOutputError(f"schema validation failed: {exc}") from exc


def schema_instruction(schema_hint: str) -> str:
    """Layer 2: prompt-side schema text."""
    return (
        "Respond with ONLY a valid JSON object matching this schema — "
        "no markdown fences, no prose outside the JSON. "
        f"Schema: {schema_hint}"
    )


async def structured_completion(
    provider: LLMProvider,
    messages: list[Message],
    *,
    model: str,
    schema: type[T],
    schema_hint: str,
    retry_feedback: str | None = None,
) -> T:
    """Full 4-layer pipeline. One bounded retry (layer 4), then raise."""
    # Build request: append schema instruction to the last user message (layer 2).
    request = list(messages)
    last_user = next((m for m in reversed(request) if m.role == "user"), None)
    if last_user is not None:
        request = [
            Message(role=m.role, content=(m.content + "\n\n" + schema_instruction(schema_hint)))
            if m is last_user else m
            for m in request
        ]
    response_format = {"type": "json_object"}
    raw = await provider.chat(request, model=model, response_format=response_format)
    try:
        return parse_structured(raw, schema)  # layers 1+2+3
    except StructuredOutputError as exc:
        # Layer 4: single bounded retry with error feedback
        retry_prompt = (
            f"Your previous response was invalid: {exc}. "
            f"Return ONLY the corrected JSON matching: {schema_hint}"
        )
        request2 = list(messages)
        request2.append(Message(role="user", content=retry_prompt))
        raw2 = await provider.chat(request2, model=model, response_format=response_format)
        try:
            return parse_structured(raw2, schema)
        except StructuredOutputError as exc2:
            raise StructuredOutputError(
                f"structured output failed after retry: {exc2}"
            ) from exc2