"""
聚合函数处理模块

功能：处理统计类查询的聚合函数提取和结果检查
"""

import re
from typing import Dict, List, Any


def extract_aggregate_functions(cypher: str) -> Dict[str, str]:
    """
    从Cypher查询中提取聚合函数及其别名
    """
    if not cypher:
        return {}
    
    aggregate_map = {}
    
    # 匹配RETURN子句（不区分大小写）
    # 匹配到 ORDER BY, SKIP, LIMIT 或字符串结尾
    return_match = re.search(r'\bRETURN\s+(.+?)(?:\s+ORDER\s+BY|\s+SKIP|\s+LIMIT|$)', cypher, re.IGNORECASE | re.DOTALL)
    if not return_match:
        return {}
    
    return_clause = return_match.group(1).strip()
    
    # 支持的聚合函数（不区分大小写）
    aggregate_functions = ['COUNT', 'SUM', 'AVG', 'MAX', 'MIN', 'COLLECT', 'COLLECT_DISTINCT']
    
    # 匹配聚合函数模式，支持以下格式：
    # - COUNT(n) AS count
    # - COUNT(n) AS `count`
    # - COUNT(n) count
    # - COUNT(DISTINCT n) AS count
    # - SUM(n.value) AS total
    # - COUNT(n) (没有别名，使用函数名)
    
    # 先尝试匹配带AS的格式：COUNT(...) AS alias
    as_pattern = r'\b(' + '|'.join(aggregate_functions) + r')\s*\([^)]*\)\s+AS\s+([`"]?)(\w+)\2'
    as_matches = re.finditer(as_pattern, return_clause, re.IGNORECASE)
    for match in as_matches:
        func_name = match.group(1).upper()
        alias = match.group(3)
        aggregate_map[alias.lower()] = func_name
    
    # 再尝试匹配不带AS但后面跟别名的格式：COUNT(...) alias
    # 需要确保alias不是下一个聚合函数
    no_as_pattern = r'\b(' + '|'.join(aggregate_functions) + r')\s*\([^)]*\)\s+([`"]?)(\w+)\2(?!\s*\()'
    no_as_matches = re.finditer(no_as_pattern, return_clause, re.IGNORECASE)
    for match in no_as_matches:
        func_name = match.group(1).upper()
        alias = match.group(3)
        # 检查alias不是聚合函数名
        if alias.upper() not in aggregate_functions:
            aggregate_map[alias.lower()] = func_name
    
    # 如果还没有找到，尝试匹配没有别名的格式：COUNT(...)
    # 这种情况下，使用函数名本身作为别名
    if not aggregate_map:
        simple_pattern = r'\b(' + '|'.join(aggregate_functions) + r')\s*\([^)]*\)'
        simple_matches = re.finditer(simple_pattern, return_clause, re.IGNORECASE)
        for match in simple_matches:
            func_name = match.group(1).upper()
            # 检查后续是否有AS子句
            after_func = return_clause[match.end():]
            as_match = re.search(r'\bAS\s+([`"]?)(\w+)\1', after_func, re.IGNORECASE)
            if as_match:
                alias = as_match.group(2)
                aggregate_map[alias.lower()] = func_name
            else:
                # 如果没有AS，使用函数名作为别名
                aggregate_map[func_name.lower()] = func_name
    
    return aggregate_map


def is_aggregate_result_zero(aggregate_results: List[Dict[str, Any]], cypher: str = None) -> bool:
    """
    检查统计类查询结果是否为0
    """
    if not aggregate_results:
        return True
    
    # 检查第一条记录的所有值
    first_result = aggregate_results[0]
    if not isinstance(first_result, dict):
        return False
    
    # 从Cypher中提取聚合函数信息
    aggregate_map = {}
    if cypher:
        aggregate_map = extract_aggregate_functions(cypher)
    
    # 如果没有提取到聚合函数信息，使用通用检查（兼容旧逻辑）
    if not aggregate_map:
        # 通用检查：任何值为None或0都认为是0结果
        for key, value in first_result.items():
            if value is None:
                return True
            if isinstance(value, (int, float)) and value == 0:
                return True
            if isinstance(value, str) and value.strip() == "0":
                return True
            if isinstance(value, list) and len(value) == 0:
                return True
        return False
    
    # 基于提取的聚合函数进行检查
    has_zero_aggregate = False
    has_non_zero_value = False
    
    for key, value in first_result.items():
        key_lower = key.lower()
        
        # 检查该字段是否对应某个聚合函数
        func_name = aggregate_map.get(key_lower)
        
        if func_name:
            # 根据聚合函数类型进行判断
            if func_name == 'COUNT':
                # COUNT结果为0表示没有数据
                if value is None or (isinstance(value, (int, float)) and value == 0):
                    has_zero_aggregate = True
                elif isinstance(value, (int, float)) and value > 0:
                    has_non_zero_value = True
            
            elif func_name == 'SUM':
                # SUM结果为0或null表示没有数据
                if value is None or (isinstance(value, (int, float)) and value == 0):
                    has_zero_aggregate = True
                elif isinstance(value, (int, float)) and value != 0:
                    has_non_zero_value = True
            
            elif func_name in ['AVG', 'MAX', 'MIN']:
                # AVG/MAX/MIN结果为null表示没有数据
                if value is None:
                    has_zero_aggregate = True
                elif isinstance(value, (int, float)) and value != 0:
                    has_non_zero_value = True
                # 注意：AVG为0可能是有效值，但通常null更常见
            
            elif func_name in ['COLLECT', 'COLLECT_DISTINCT']:
                # COLLECT结果为空列表表示没有数据
                if value is None or (isinstance(value, list) and len(value) == 0):
                    has_zero_aggregate = True
                elif isinstance(value, list) and len(value) > 0:
                    has_non_zero_value = True
        else:
            # 如果字段不在聚合函数映射中，但值不为0/null，说明可能有其他有效数据
            if value is not None:
                if isinstance(value, (int, float)) and value != 0:
                    has_non_zero_value = True
                elif isinstance(value, str) and value.strip() and value.strip() != "0":
                    has_non_zero_value = True
                elif isinstance(value, list) and len(value) > 0:
                    has_non_zero_value = True
    
    # 如果有聚合函数结果为0/null，且没有其他非0值，认为是0结果
    if has_zero_aggregate and not has_non_zero_value:
        return True
    
    # 如果所有聚合函数值都不为0或null，返回False
    return False

