"""Parse LLM text responses into typed Pydantic models."""

from __future__ import annotations

import re
from typing import TypeVar

import orjson
from pydantic import BaseModel, ValidationError

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

# Matches ```json ... ``` or ``` ... ``` code fences
_CODE_FENCE_RE = re.compile(r"```(?:json)?\s*\n?(.*?)\n?\s*```", re.DOTALL)

# Matches the outermost { ... } or [ ... ]
_JSON_OBJECT_RE = re.compile(r"(\{.*\}|\[.*\])", re.DOTALL)


class OutputParseError(Exception):
    """Raised when an LLM response cannot be parsed into the target model."""

    def __init__(self, message: str, raw_text: str, cause: Exception | None = None) -> None:
        self.raw_text = raw_text
        self.cause = cause
        super().__init__(message)


def _extract_json(text: str) -> str:
    """Extract JSON from text that may contain markdown fences or surrounding prose."""
    # Try code fence first
    fence_match = _CODE_FENCE_RE.search(text)
    if fence_match:
        return fence_match.group(1).strip()

    # Try raw JSON object/array
    obj_match = _JSON_OBJECT_RE.search(text)
    if obj_match:
        return obj_match.group(1).strip()

    # Return as-is and let the JSON parser handle the error
    return text.strip()


def parse_response(text: str, model_class: type[T]) -> T:
    """Parse an LLM response string into a Pydantic model instance.

    Handles:
    - Pure JSON responses
    - JSON wrapped in markdown code fences (```json ... ```)
    - JSON embedded in surrounding prose

    Args:
        text: Raw LLM response text.
        model_class: Pydantic BaseModel subclass to parse into.

    Returns:
        Validated instance of model_class.

    Raises:
        OutputParseError: If JSON extraction or Pydantic validation fails.
    """
    json_str = _extract_json(text)

    try:
        data = orjson.loads(json_str)
    except (orjson.JSONDecodeError, ValueError) as e:
        raise OutputParseError(
            f"Failed to parse JSON from LLM response: {e}\n"
            f"Extracted text: {json_str[:200]}{'...' if len(json_str) > 200 else ''}",
            raw_text=text,
            cause=e,
        ) from e

    try:
        return model_class.model_validate(data)
    except ValidationError as e:
        raise OutputParseError(
            f"JSON parsed but failed {model_class.__name__} validation:\n{e}",
            raw_text=text,
            cause=e,
        ) from e
