"""Safe JSON parsing utilities for LLM response text."""

from __future__ import annotations

import json
from typing import Any, Literal


class LLMParseError(ValueError):
    """Structured error for LLM JSON parsing failures.

    Attributes:
        raw_text: The original text that failed to parse (truncated).
        reason: Human-readable failure reason.
    """

    def __init__(self, reason: str, raw_text: str = ""):
        self.reason = reason
        self.raw_text = raw_text[:200] if raw_text else ""
        super().__init__(f"LLM JSON parse error: {reason}")


def _strip_markdown_fences(text: str) -> str:
    """Remove markdown code fences from LLM output.

    Handles: ```json, ```, ```python, etc.
    """
    text = text.strip()
    if text.startswith("```"):
        lines = text.split("\n")
        # Remove opening fence (first line starting with ```)
        if lines and lines[0].strip().startswith("```"):
            lines = lines[1:]
        # Remove closing fence (last line that is just ```)
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        text = "\n".join(lines).strip()
    return text


def parse_llm_json(
    text: str,
    *,
    expect: Literal["object", "array", "any"] = "any",
) -> dict[str, Any] | list[Any]:
    """Parse JSON from LLM response text with markdown fence stripping.

    Handles common LLM output patterns:
    - Clean JSON
    - JSON wrapped in markdown code fences (```json ... ```)
    - JSON with surrounding explanation text
    - JSON arrays or objects

    Args:
        text: Raw LLM response text.
        expect: Expected JSON type. "object" requires dict, "array" requires list.

    Returns:
        Parsed JSON (dict or list).

    Raises:
        LLMParseError: If parsing fails or type doesn't match expectation.
    """
    if not text or not text.strip():
        raise LLMParseError("Empty response text", text or "")

    cleaned = _strip_markdown_fences(text)

    # Try direct parse first
    try:
        result = json.loads(cleaned)
        return _validate_type(result, expect, text)
    except json.JSONDecodeError:
        pass

    # Try to find JSON object or array in text — pick earliest match
    candidates = []
    for start_char in ("{", "["):
        pos = cleaned.find(start_char)
        if pos != -1:
            candidates.append(pos)

    # Sort by position so we try the earliest JSON structure first
    candidates.sort()
    for start in candidates:
        for end in range(len(cleaned), start, -1):
            try:
                result = json.loads(cleaned[start:end])
                return _validate_type(result, expect, text)
            except json.JSONDecodeError:
                continue

    raise LLMParseError("No valid JSON found in response", text)


def _validate_type(
    result: Any,
    expect: Literal["object", "array", "any"],
    raw_text: str,
) -> dict[str, Any] | list[Any]:
    """Validate that parsed JSON matches the expected type."""
    if expect == "object" and not isinstance(result, dict):
        raise LLMParseError(f"Expected JSON object, got {type(result).__name__}", raw_text)
    if expect == "array" and not isinstance(result, list):
        raise LLMParseError(f"Expected JSON array, got {type(result).__name__}", raw_text)
    return result


def parse_json_safe(text: str) -> dict | list | None:
    """Parse JSON text, returning None on failure instead of raising."""
    try:
        return json.loads(text)
    except (json.JSONDecodeError, TypeError, ValueError):
        return None