"""
JSON工具类，支持复杂的键路径提取功能
支持点号分割、通配符、正则表达式和区间匹配
"""

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

try:
    import numpy as np
    HAS_NUMPY = True
except ImportError:
    HAS_NUMPY = False


class JsonPathExtractor:
    """JSON路径提取器，支持多种路径模式"""

    @staticmethod
    def extract_by_path(data: Any, key_path: str) -> Any:
        """
        根据键路径从JSON数据中提取值

        支持的路径格式：
        - 普通路径：'key1.key2.key3'
        - 带引号的键：'key1."key with space".key3'
        - 通配符：'key1.*.key3' (返回数组)
        - 正则表达式：'key1.r"^[0-3]$".key3' (返回数组)
        - 区间匹配：'key1.[0:10].key3' (返回数组)
        """
        if not key_path or data is None:
            return None

        segments = JsonPathExtractor._parse_path(key_path)
        if not segments:
            return None

        return JsonPathExtractor._extract_recursive(data, segments)

    @staticmethod
    def _parse_path(key_path: str) -> List[dict]:
        """解析路径字符串为segments列表"""
        segments = []
        current_segment = ""
        in_quotes = False
        quote_char = None
        i = 0

        while i < len(key_path):
            char = key_path[i]

            if char in ['"', "'"]:
                if not in_quotes:
                    in_quotes = True
                    quote_char = char
                elif char == quote_char:
                    in_quotes = False
                    quote_char = None
                current_segment += char
            elif char == '.' and not in_quotes:
                if current_segment:
                    segment = JsonPathExtractor._parse_segment(current_segment)
                    if segment:
                        segments.append(segment)
                current_segment = ""
            else:
                current_segment += char

            i += 1

        if current_segment:
            segment = JsonPathExtractor._parse_segment(current_segment)
            if segment:
                segments.append(segment)

        return segments

    @staticmethod
    def _parse_segment(segment: str) -> Optional[dict]:
        """解析单个路径段"""
        segment = segment.strip()
        if not segment:
            return None

        # 区间匹配：[start:end]
        range_match = re.match(r'^\[(\d+):(\d+)\]$', segment)
        if range_match:
            start = int(range_match.group(1))
            end = int(range_match.group(2))
            return {'type': 'range', 'start': start, 'end': end}

        # 正则表达式：r"pattern"
        regex_match = re.match(r'^r["\'](.+)["\']$', segment)
        if regex_match:
            pattern = regex_match.group(1)
            try:
                compiled_pattern = re.compile(pattern)
                return {'type': 'regex', 'pattern': compiled_pattern}
            except re.error:
                return {'type': 'key', 'value': segment}

        # 通配符：*
        if segment == '*':
            return {'type': 'wildcard'}

        # 带引号的键："key name"
        quoted_match = re.match(r'^["\'](.+)["\']$', segment)
        if quoted_match:
            return {'type': 'key', 'value': quoted_match.group(1)}

        # 普通键
        return {'type': 'key', 'value': segment}

    @staticmethod
    def _extract_recursive(data: Any, segments: List[dict]) -> Any:
        """递归提取数据"""
        if not segments:
            return data

        current_segment = segments[0]
        remaining_segments = segments[1:]
        segment_type = current_segment['type']

        if segment_type == 'key':
            key = current_segment['value']
            if isinstance(data, dict) and key in data:
                return JsonPathExtractor._extract_recursive(data[key], remaining_segments)
            elif isinstance(data, (list, tuple)):
                try:
                    index = int(key)
                    if 0 <= index < len(data):
                        return JsonPathExtractor._extract_recursive(data[index], remaining_segments)
                except (ValueError, IndexError):
                    pass
            return None

        elif segment_type == 'wildcard':
            return JsonPathExtractor._extract_wildcard(data, remaining_segments)

        elif segment_type == 'regex':
            pattern = current_segment['pattern']
            return JsonPathExtractor._extract_regex(data, pattern, remaining_segments)

        elif segment_type == 'range':
            start = current_segment['start']
            end = current_segment['end']
            return JsonPathExtractor._extract_range(data, start, end, remaining_segments)

        return None

    @staticmethod
    def _extract_wildcard(data: Any, remaining_segments: List[dict]) -> List[Any]:
        """通配符提取"""
        results = []

        if isinstance(data, dict):
            for value in data.values():
                result = JsonPathExtractor._extract_recursive(value, remaining_segments)
                if result is not None:
                    results.append(result)
        elif isinstance(data, (list, tuple)):
            for item in data:
                result = JsonPathExtractor._extract_recursive(item, remaining_segments)
                if result is not None:
                    results.append(result)

        return results if results else None

    @staticmethod
    def _extract_regex(data: Any, pattern: re.Pattern, remaining_segments: List[dict]) -> List[Any]:
        """正则表达式提取"""
        results = None
        if isinstance(data, dict):
            results = {}
            for key, value in data.items():
                if pattern.match(str(key)):
                    result = JsonPathExtractor._extract_recursive(value, remaining_segments)
                    if result is not None:
                        results[key] = result
        elif isinstance(data, (list, tuple)):
            results = []
            for index, item in enumerate(data):
                if pattern.match(str(index)):
                    result = JsonPathExtractor._extract_recursive(item, remaining_segments)
                    if result is not None:
                        results.append(result)

        return results if results else None

    @staticmethod
    def _extract_range(data: Any, start: int, end: int, remaining_segments: List[dict]) -> List[Any]:
        """区间提取"""
        if not isinstance(data, (list, tuple)):
            return None

        results = []
        start = max(0, start)
        end = min(len(data), end)

        if start >= end:
            return None

        for i in range(start, end):
            result = JsonPathExtractor._extract_recursive(data[i], remaining_segments)
            if result is not None:
                results.append(result)

        return results if results else None


def extract_json_by_path(data: Any, key_path: str) -> Any:
    """
    便捷函数：根据键路径从JSON数据中提取值

    Examples:
        extract_json_by_path({"a": {"b": 1}}, "a.b")  # 返回: 1
        extract_json_by_path({"a": {"x": 1, "y": 2}}, "a.*")  # 返回: [1, 2]
        extract_json_by_path({"a": [1, 2, 3, 4, 5]}, "a.[1:4]")  # 返回: [2, 3, 4]
    """
    return JsonPathExtractor.extract_by_path(data, key_path)


class NumpyJSONEncoder(json.JSONEncoder):
    """支持numpy数据类型的JSON编码器"""

    def default(self, obj):
        if not HAS_NUMPY:
            return super().default(obj)

        if isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.bool_):
            return bool(obj)
        elif isinstance(obj, np.complexfloating):
            return {"real": float(obj.real), "imag": float(obj.imag)}

        return super().default(obj)


def json_dumps_numpy(obj, **kwargs) -> str:
    """支持numpy数据类型的json.dumps函数"""
    if 'cls' not in kwargs:
        kwargs['cls'] = NumpyJSONEncoder
    return json.dumps(obj, **kwargs)


def json_dumps_safe(obj, **kwargs) -> str:
    """安全的json.dumps函数，自动处理numpy类型和其他常见的序列化问题"""
    defaults = {
        'ensure_ascii': False,
        'cls': NumpyJSONEncoder
    }
    for key, value in defaults.items():
        if key not in kwargs:
            kwargs[key] = value
    return json.dumps(obj, **kwargs)
