"""
JSON 安全解析功能组件

提供安全的 JSON 解析功能，专门处理大模型返回的包装 JSON 格式
"""

import json
import re
from typing import Any, Dict, Optional, Union, List
from .base import SimpleFunction, component_function


class JSONParseError(Exception):
    """JSON 解析异常"""
    pass


class JsonSafeParseFunction(SimpleFunction[str, Any]):
    """JSON 安全解析函数组件
    
    专门处理大模型返回的包装 JSON，支持：
    1. 去除 ```json ``` 包装格式
    2. 安全的 JSON 解析
    3. 严格的异常处理
    """
    
    def __init__(self):
        """初始化 JSON 安全解析函数"""
        super().__init__(
            name="json_safe_parse",
            description="安全解析 JSON 字符串，支持处理大模型返回的包装格式",
            func=self._parse_json
        )
    
    def _parse_json(self, input_data: str, context: Optional[Dict[str, Any]] = None) -> Any:
        """执行 JSON 解析
        
        Args:
            input_data: 待解析的 JSON 字符串
            context: 上下文信息，可包含解析选项
            
        Returns:
            解析后的 Python 对象
            
        Raises:
            JSONParseError: 当 JSON 解析失败时抛出
        """
        if not isinstance(input_data, str):
            raise JSONParseError(f"Expected string input, got {type(input_data)}")
        
        # 获取上下文选项
        options = context.get('options', {}) if context else {}
        strict_mode = options.get('strict', True)  # 默认严格模式
        
        try:
            # 1. 清理输入字符串
            cleaned_json = self._extract_json_from_wrapped_string(input_data.strip())
            
            # 2. 尝试解析 JSON
            if not cleaned_json:
                raise JSONParseError("Empty or invalid JSON content after extraction")
            
            parsed_data = json.loads(cleaned_json, strict=strict_mode)
            return parsed_data
            
        except json.JSONDecodeError as e:
            raise JSONParseError(f"JSON decode error: {str(e)}")
        except Exception as e:
            raise JSONParseError(f"Unexpected error during JSON parsing: {str(e)}")
    
    def _extract_json_from_wrapped_string(self, text: str) -> str:
        """从包装的字符串中提取 JSON 内容
        
        支持的格式：
        - ```json { ... } ```
        - ```{ ... }```
        - { ... }
        - [ ... ]
        
        Args:
            text: 包装的字符串
            
        Returns:
            提取的 JSON 字符串
        """
        # 去除首尾空白
        text = text.strip()
        
        # 模式1: ```json ... ``` 或 ```... ```
        json_block_pattern = r'```(?:json)?\s*\n?(.*?)\n?```'
        match = re.search(json_block_pattern, text, re.DOTALL | re.IGNORECASE)
        if match:
            return match.group(1).strip()
        
        # 模式2: 直接的 JSON 对象或数组
        # 寻找完整的 JSON 结构
        json_patterns = [
            r'(\{.*\})',  # JSON 对象
            r'(\[.*\])',  # JSON 数组
        ]
        
        for pattern in json_patterns:
            match = re.search(pattern, text, re.DOTALL)
            if match:
                potential_json = match.group(1).strip()
                # 简单验证是否是有效的 JSON 开始
                if (potential_json.startswith('{') and potential_json.endswith('}')) or \
                   (potential_json.startswith('[') and potential_json.endswith(']')):
                    return potential_json
        
        # 如果没有找到明确的 JSON 结构，返回原文本
        return text


# 创建全局实例
json_safe_parse_function = JsonSafeParseFunction()


@component_function(
    name="json_safe_parse",
    description="安全解析 JSON 字符串，支持处理大模型返回的包装格式",
    category="json_utils",
    auto_register=True
)
def safe_json_parse(input_data: str, context: Optional[Dict[str, Any]] = None) -> Any:
    """安全解析 JSON 字符串的便捷函数
    
    Args:
        input_data: 待解析的 JSON 字符串
        context: 上下文信息
        
    Returns:
        解析后的 Python 对象
        
    Raises:
        JSONParseError: 当 JSON 解析失败时抛出
    """
    return json_safe_parse_function(input_data, context)


@component_function(
    name="batch_json_parse",
    description="批量解析 JSON 字符串列表",
    category="json_utils",
    auto_register=True
)
def batch_json_parse(input_data: List[str], context: Optional[Dict[str, Any]] = None) -> List[Any]:
    """批量解析 JSON 字符串
    
    Args:
        input_data: 待解析的 JSON 字符串列表
        context: 上下文信息
        
    Returns:
        解析后的 Python 对象列表
        
    Raises:
        JSONParseError: 当任何 JSON 解析失败时抛出
    """
    if not isinstance(input_data, list):
        raise JSONParseError(f"Expected list input, got {type(input_data)}")
    
    results = []
    for i, json_str in enumerate(input_data):
        try:
            result = safe_json_parse(json_str, context)
            results.append(result)
        except JSONParseError as e:
            raise JSONParseError(f"Failed to parse JSON at index {i}: {str(e)}")
    
    return results


@component_function(
    name="validate_json_structure",
    description="验证 JSON 结构是否符合预期",
    category="json_utils",
    auto_register=True
)
def validate_json_structure(
    input_data: Union[str, Dict, List], 
    context: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """验证 JSON 结构
    
    Args:
        input_data: JSON 字符串或已解析的对象
        context: 上下文信息，可包含验证规则
        
    Returns:
        验证结果字典，包含 'valid', 'data', 'errors' 字段
    """
    result = {
        'valid': False,
        'data': None,
        'errors': []
    }
    
    try:
        # 如果是字符串，先解析
        if isinstance(input_data, str):
            data = safe_json_parse(input_data, context)
        else:
            data = input_data
        
        result['data'] = data
        
        # 获取验证规则
        validation_rules = context.get('validation_rules', {}) if context else {}
        
        # 基本验证
        if not isinstance(data, (dict, list)):
            result['errors'].append("Parsed data is not a JSON object or array")
            return result
        
        # 自定义验证规则
        if 'required_keys' in validation_rules and isinstance(data, dict):
            required_keys = validation_rules['required_keys']
            missing_keys = [key for key in required_keys if key not in data]
            if missing_keys:
                result['errors'].append(f"Missing required keys: {missing_keys}")
        
        if 'expected_type' in validation_rules:
            expected_type = validation_rules['expected_type']
            if expected_type == 'object' and not isinstance(data, dict):
                result['errors'].append("Expected JSON object")
            elif expected_type == 'array' and not isinstance(data, list):
                result['errors'].append("Expected JSON array")
        
        # 如果没有错误，标记为有效
        result['valid'] = len(result['errors']) == 0
        
    except JSONParseError as e:
        result['errors'].append(str(e))
    except Exception as e:
        result['errors'].append(f"Validation error: {str(e)}")
    
    return result