#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
JSON提取和解析工具类

功能: 从LLM响应中提取和解析JSON内容
"""

import json
import re
from typing import Any

from src.utils.logger import get_logger

logger = get_logger(__name__)


class JSONExtractor:
    """JSON提取和解析工具类"""

    @staticmethod
    def extract_json_from_response(content: str) -> str:
        """从LLM响应中提取JSON内容

        Args:
            content: LLM返回的原始文本内容

        Returns:
            提取出的JSON字符串
        """
        content = content.strip()
        if "```json" in content:
            content = content.split("```json")[1].split("```")[0].strip()
        elif "```" in content:
            content = content.split("```")[1].split("```")[0].strip()
        return content

    @staticmethod
    def safe_parse_json(content: str, default: Any = None) -> Any:
        """安全解析JSON,失败时返回默认值

        Args:
            content: JSON字符串
            default: 解析失败时返回的默认值

        Returns:
            解析后的Python对象,失败时返回default
        """
        try:
            return json.loads(content)
        except json.JSONDecodeError as e:
            # 尝试修复常见的JSON格式问题
            try:
                fixed_content = JSONExtractor._fix_json_format(content)
                return json.loads(fixed_content)
            except Exception as fix_error:
                logger.error(f"JSON解析失败: {e}\nJSON修复也失败: {fix_error}\n内容: {content}...")
                return default if default is not None else {}

    @staticmethod
    def _fix_json_format(content: str) -> str:
        """修复常见的JSON格式问题

        Args:
            content: 可能有格式问题的JSON字符串

        Returns:
            修复后的JSON字符串
        """
        # 移除或替换控制字符（保留合法的 \n \t \r）
        # 先将合法的转义字符临时替换
        content = content.replace('\\n', '\x00NEWLINE\x00')
        content = content.replace('\\t', '\x00TAB\x00')
        content = content.replace('\\r', '\x00RETURN\x00')

        # 移除所有ASCII控制字符（0-31，除了合法的换行、制表、回车）
        content = re.sub(r'[\x00-\x08\x0B\x0C\x0E-\x1F]', '', content)

        # 恢复合法的转义字符
        content = content.replace('\x00NEWLINE\x00', '\\n')
        content = content.replace('\x00TAB\x00', '\\t')
        content = content.replace('\x00RETURN\x00', '\\r')

        # 移除对象中的尾随逗号: ,} 或 , } 或 ,\n}
        content = re.sub(r',\s*}', '}', content)

        # 移除数组中的尾随逗号: ,] 或 , ] 或 ,\n]
        content = re.sub(r',\s*]', ']', content)

        return content

    @staticmethod
    def extract_and_parse(content: str, default: Any = None) -> Any:
        """一步完成提取和解析

        Args:
            content: LLM返回的原始文本内容
            default: 解析失败时返回的默认值

        Returns:
            解析后的Python对象,失败时返回default
        """
        json_str = JSONExtractor.extract_json_from_response(content)
        return JSONExtractor.safe_parse_json(json_str, default)
