#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
JSON Processor
JSON处理器 - 使用json_repair处理LLM输出，修复格式问题
"""

import json
import logging
import re
from typing import Any, Dict, List, Optional, Union

try:
    from json_repair import repair_json
    JSON_REPAIR_AVAILABLE = True
except ImportError:
    JSON_REPAIR_AVAILABLE = False
    
from ..exceptions import ValidationError

logger = logging.getLogger(__name__)


class JSONProcessor:
    """
    JSON处理器
    专门处理LLM输出的JSON数据，包括格式修复和验证
    """

    def __init__(self, enable_repair: bool = True, max_repair_attempts: int = 3):
        """
        初始化JSON处理器

        Args:
            enable_repair: 是否启用JSON修复
            max_repair_attempts: 最大修复尝试次数
        """
        self.enable_repair = enable_repair and JSON_REPAIR_AVAILABLE
        self.max_repair_attempts = max_repair_attempts
        
        if not JSON_REPAIR_AVAILABLE and enable_repair:
            logger.warning("json_repair不可用，将禁用JSON修复功能")
        
        logger.debug(f"JSON处理器初始化完成，修复功能: {self.enable_repair}")

    def parse_llm_output(
        self, 
        raw_output: str, 
        expected_fields: Optional[List[str]] = None,
        required_fields: Optional[List[str]] = None
    ) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
        """
        解析LLM输出的JSON数据

        Args:
            raw_output: LLM原始输出
            expected_fields: 期望的字段列表
            required_fields: 必需的字段列表

        Returns:
            解析后的JSON数据

        Raises:
            ValidationError: 解析失败时抛出
        """
        if not raw_output or not raw_output.strip():
            raise ValidationError("输入内容为空")

        # 尝试多种解析方法
        parsed_data = (
            self._parse_json_block(raw_output) or
            self._parse_direct_json(raw_output) or
            self._parse_with_repair(raw_output) or
            self._parse_structured_text(raw_output)
        )

        if parsed_data is None:
            raise ValidationError("无法解析JSON数据")

        # 验证数据结构
        if expected_fields or required_fields:
            self._validate_fields(parsed_data, expected_fields, required_fields)

        return parsed_data

    def _parse_json_block(self, content: str) -> Optional[Union[Dict, List]]:
        """解析```json```代码块中的JSON"""
        try:
            # 匹配 ```json 代码块
            json_pattern = re.compile(r'```json\s*(.*?)\s*```', re.DOTALL | re.IGNORECASE)
            matches = json_pattern.findall(content)
            
            for match in matches:
                json_str = match.strip()
                if json_str:
                    try:
                        return json.loads(json_str)
                    except json.JSONDecodeError:
                        if self.enable_repair:
                            repaired = self._repair_json_string(json_str)
                            if repaired:
                                return repaired
            return None
        except Exception as e:
            logger.debug(f"JSON代码块解析失败: {e}")
            return None

    def _parse_direct_json(self, content: str) -> Optional[Union[Dict, List]]:
        """直接解析JSON内容"""
        try:
            # 尝试直接解析整个内容
            content_cleaned = content.strip()
            return json.loads(content_cleaned)
        except json.JSONDecodeError:
            # 尝试提取JSON部分
            json_patterns = [
                re.compile(r'\{.*\}', re.DOTALL),  # 匹配JSON对象
                re.compile(r'\[.*\]', re.DOTALL),  # 匹配JSON数组
            ]
            
            for pattern in json_patterns:
                matches = pattern.findall(content)
                for match in matches:
                    try:
                        return json.loads(match)
                    except json.JSONDecodeError:
                        continue
            return None
        except Exception as e:
            logger.debug(f"直接JSON解析失败: {e}")
            return None

    def _parse_with_repair(self, content: str) -> Optional[Union[Dict, List]]:
        """使用json_repair修复并解析JSON"""
        if not self.enable_repair:
            return None

        try:
            # 尝试修复整个内容
            repaired = self._repair_json_string(content)
            if repaired:
                return repaired

            # 尝试提取可能的JSON部分进行修复
            potential_json_patterns = [
                r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}',  # 嵌套JSON对象
                r'\[[^\[\]]*(?:\[[^\[\]]*\][^\[\]]*)*\]',  # 嵌套JSON数组
            ]

            for pattern in potential_json_patterns:
                matches = re.findall(pattern, content, re.DOTALL)
                for match in matches:
                    repaired = self._repair_json_string(match)
                    if repaired:
                        return repaired

            return None
        except Exception as e:
            logger.debug(f"JSON修复解析失败: {e}")
            return None

    def _repair_json_string(self, json_str: str) -> Optional[Union[Dict, List]]:
        """修复JSON字符串"""
        if not self.enable_repair:
            return None

        for attempt in range(self.max_repair_attempts):
            try:
                # 使用json_repair修复
                repaired_str = repair_json(json_str)
                parsed = json.loads(repaired_str)
                
                logger.debug(f"JSON修复成功，尝试次数: {attempt + 1}")
                return parsed
                
            except Exception as e:
                logger.debug(f"JSON修复尝试 {attempt + 1} 失败: {e}")
                # 尝试一些简单的预处理
                json_str = self._preprocess_json_string(json_str)

        return None

    def _preprocess_json_string(self, json_str: str) -> str:
        """预处理JSON字符串"""
        # 移除常见的问题字符
        processed = json_str.strip()
        
        # 移除开头和结尾的多余字符
        processed = re.sub(r'^[^{\[]*', '', processed)
        processed = re.sub(r'[^}\]]*$', '', processed)
        
        # 修复常见的引号问题
        processed = re.sub(r'([{,]\s*)([a-zA-Z_][a-zA-Z0-9_]*)\s*:', r'\1"\2":', processed)
        
        # 修复尾随逗号
        processed = re.sub(r',\s*([}\]])', r'\1', processed)
        
        return processed

    def _parse_structured_text(self, content: str) -> Optional[List[Dict[str, Any]]]:
        """解析结构化文本格式（回退方案）"""
        try:
            # 这里可以实现特定的结构化文本解析逻辑
            # 例如解析类似markdown或特定格式的文本
            
            # 简单示例：解析键值对格式
            result = []
            current_item = {}
            
            lines = content.split('\n')
            for line in lines:
                line = line.strip()
                if not line:
                    if current_item:
                        result.append(current_item)
                        current_item = {}
                    continue
                
                if ':' in line:
                    key, value = line.split(':', 1)
                    key = key.strip().lower()
                    value = value.strip()
                    current_item[key] = value
            
            if current_item:
                result.append(current_item)
            
            return result if result else None
            
        except Exception as e:
            logger.debug(f"结构化文本解析失败: {e}")
            return None

    def _validate_fields(
        self, 
        data: Union[Dict, List], 
        expected_fields: Optional[List[str]] = None,
        required_fields: Optional[List[str]] = None
    ) -> None:
        """验证字段完整性"""
        
        def validate_item(item: Dict[str, Any]) -> None:
            if not isinstance(item, dict):
                raise ValidationError(f"期望字典类型，实际类型: {type(item)}")
            
            # 检查必需字段
            if required_fields:
                missing_fields = [field for field in required_fields if field not in item]
                if missing_fields:
                    raise ValidationError(f"缺少必需字段: {missing_fields}")
            
            # 记录意外字段
            if expected_fields:
                unexpected_fields = [field for field in item.keys() if field not in expected_fields]
                if unexpected_fields:
                    logger.debug(f"发现意外字段: {unexpected_fields}")

        if isinstance(data, list):
            for i, item in enumerate(data):
                try:
                    validate_item(item)
                except ValidationError as e:
                    raise ValidationError(f"数组项 {i}: {e}")
        elif isinstance(data, dict):
            validate_item(data)
        else:
            raise ValidationError(f"不支持的数据类型: {type(data)}")

    def format_output(
        self, 
        data: Union[Dict, List], 
        output_format: str = "json"
    ) -> str:
        """
        格式化输出数据

        Args:
            data: 要格式化的数据
            output_format: 输出格式 (json, pretty_json)

        Returns:
            格式化后的字符串
        """
        try:
            if output_format == "pretty_json":
                return json.dumps(data, ensure_ascii=False, indent=2)
            else:
                return json.dumps(data, ensure_ascii=False)
        except Exception as e:
            logger.error(f"数据格式化失败: {e}")
            return str(data)

    def validate_json_schema(
        self, 
        data: Union[Dict, List], 
        schema: Dict[str, Any]
    ) -> bool:
        """
        验证JSON模式（简单版本）

        Args:
            data: 要验证的数据
            schema: JSON模式定义

        Returns:
            验证是否通过
        """
        # 这里可以实现更复杂的JSON模式验证
        # 目前提供基础版本
        
        try:
            if "type" in schema:
                expected_type = schema["type"]
                if expected_type == "object" and not isinstance(data, dict):
                    return False
                elif expected_type == "array" and not isinstance(data, list):
                    return False
            
            if "required" in schema and isinstance(data, dict):
                required_fields = schema["required"]
                for field in required_fields:
                    if field not in data:
                        return False
            
            return True
            
        except Exception as e:
            logger.error(f"JSON模式验证失败: {e}")
            return False 