#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
JSON Utils - JSON parsing utilities
Specifically for extracting and parsing JSON data from LLM output,
supports Markdown code block processing and is robust to Python
triple-quoted strings generated by LLMs.
"""

import json
import re
from typing import Any


def _escape_triple_quoted_strings(text: str) -> str:
    """
    Convert Python triple-quoted strings (\"\"\" ... \"\"\")
    into JSON-safe escaped strings.

    This fixes cases where LLMs generate Python code inside JSON,
    which would otherwise violate JSON syntax.
    """

    def replacer(match: re.Match) -> str:
        content = match.group(1)
        # json.dumps safely escapes newlines and quotes
        return json.dumps(content)

    pattern = re.compile(r'"""([\s\S]*?)"""')
    return pattern.sub(replacer, text)


def extract_json_from_text(text: str) -> dict[str, Any] | list[Any] | None:
    """
    Extract JSON object or array from text.

    Supports:
    1. ```json ... ``` code blocks
    2. ``` ... ``` code blocks
    3. Pure JSON text
    4. JSON containing Python triple-quoted strings (sanitized)

    Args:
        text: Original text containing JSON

    Returns:
        Parsed JSON object (dict) or array (list),
        or None if parsing fails.
    """
    if not text:
        return None

    # 🔧 FIX: sanitize triple-quoted strings before any JSON parsing
    text = _escape_triple_quoted_strings(text)

    # 1. Try matching Markdown code blocks
    code_block_pattern = re.compile(r"```(?:json)?\s*([\s\S]*?)\s*```")
    match = code_block_pattern.search(text)

    if match:
        json_str = match.group(1).strip()
        try:
            return json.loads(json_str)
        except json.JSONDecodeError:
            pass

    # 2. Try parsing the full text directly
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # 3. Try extracting outermost JSON object
    json_obj_pattern = re.compile(r"\{[\s\S]*\}")
    match_obj = json_obj_pattern.search(text)
    if match_obj:
        try:
            return json.loads(match_obj.group(0))
        except json.JSONDecodeError:
            pass

    # 4. Try extracting outermost JSON array
    json_arr_pattern = re.compile(r"\[[\s\S]*\]")
    match_arr = json_arr_pattern.search(text)
    if match_arr:
        try:
            return json.loads(match_arr.group(0))
        except json.JSONDecodeError:
            pass

    return None


def clean_json_string(json_str: str) -> str:
    """
    Clean JSON string by removing illegal control characters.
    """
    return re.sub(r"[\x00-\x1f\x7f-\x9f]", "", json_str)
