"""
数据清洗处理器

提供各种数据清洗功能，如空值处理、重复项处理、数据类型转换等
"""

import pandas as pd
from typing import Dict, Any, Optional, List, Union
from app.core.utils.log_utils import get_logger

logger = get_logger(__name__)


class DataCleaner:
    """数据清洗处理器
    
    提供标准化的数据清洗功能，支持链式调用和规则配置
    """
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """初始化数据清洗器
        
        Args:
            config: 清洗配置
        """
        self.config = config or {}
        self.cleaning_rules = []
    
    def add_rule(self, rule_type: str, **kwargs):
        """添加清洗规则
        
        Args:
            rule_type: 规则类型
            **kwargs: 规则参数
        """
        rule = {'type': rule_type, **kwargs}
        self.cleaning_rules.append(rule)
        logger.debug(f"添加清洗规则: {rule_type}")
    
    def clean(self, df: pd.DataFrame) -> pd.DataFrame:
        """执行数据清洗
        
        Args:
            df: 输入数据
            
        Returns:
            清洗后的数据
        """
        logger.info(f"开始数据清洗，原始数据形状: {df.shape}")
        
        result_df = df.copy()
        
        for i, rule in enumerate(self.cleaning_rules):
            try:
                logger.debug(f"执行清洗规则 {i+1}/{len(self.cleaning_rules)}: {rule['type']}")
                result_df = self._apply_rule(result_df, rule)
                logger.debug(f"规则执行完成，数据形状: {result_df.shape}")
            except Exception as e:
                logger.error(f"清洗规则执行失败: {rule}, 错误: {e}")
                # 继续执行下一个规则，而不是中断整个流程
                continue
        
        logger.info(f"数据清洗完成，最终数据形状: {result_df.shape}")
        return result_df
    
    def _apply_rule(self, df: pd.DataFrame, rule: Dict[str, Any]) -> pd.DataFrame:
        """应用单个清洗规则
        
        Args:
            df: 数据
            rule: 规则配置
            
        Returns:
            处理后的数据
        """
        rule_type = rule.get('type')
        
        if rule_type == 'remove_duplicates':
            return self._remove_duplicates(df, rule)
        elif rule_type == 'fill_na':
            return self._fill_na(df, rule)
        elif rule_type == 'remove_rows':
            return self._remove_rows(df, rule)
        elif rule_type == 'convert_type':
            return self._convert_type(df, rule)
        elif rule_type == 'strip_whitespace':
            return self._strip_whitespace(df, rule)
        elif rule_type == 'normalize_text':
            return self._normalize_text(df, rule)
        elif rule_type == 'validate_data':
            return self._validate_data(df, rule)
        else:
            logger.warning(f"未知的清洗规则类型: {rule_type}")
            return df
    
    def _remove_duplicates(self, df: pd.DataFrame, rule: Dict[str, Any]) -> pd.DataFrame:
        """移除重复项
        
        Args:
            df: 数据
            rule: 规则配置
            
        Returns:
            处理后的数据
        """
        subset = rule.get('subset')  # 用于判断重复的列
        keep = rule.get('keep', 'first')  # 保留哪个重复项
        
        before_count = len(df)
        df_cleaned = df.drop_duplicates(subset=subset, keep=keep)
        after_count = len(df_cleaned)
        
        logger.info(f"移除重复项: {before_count - after_count} 行被移除")
        return df_cleaned
    
    def _fill_na(self, df: pd.DataFrame, rule: Dict[str, Any]) -> pd.DataFrame:
        """填充空值
        
        Args:
            df: 数据
            rule: 规则配置
            
        Returns:
            处理后的数据
        """
        columns = rule.get('columns')  # 要处理的列
        value = rule.get('value', 0)  # 填充值
        method = rule.get('method')  # 填充方法（'ffill', 'bfill', 'mean', 'median'）
        
        if columns:
            # 处理指定列
            if isinstance(columns, str):
                columns = [columns]
            
            for col in columns:
                if col in df.columns:
                    if method == 'ffill':
                        df[col] = df[col].fillna(method='ffill')
                    elif method == 'bfill':
                        df[col] = df[col].fillna(method='bfill')
                    elif method == 'mean':
                        df[col] = df[col].fillna(df[col].mean())
                    elif method == 'median':
                        df[col] = df[col].fillna(df[col].median())
                    else:
                        df[col] = df[col].fillna(value)
                    
                    logger.debug(f"填充列 {col} 的空值: {method or value}")
        else:
            # 处理所有列
            if method == 'ffill':
                df = df.fillna(method='ffill')
            elif method == 'bfill':
                df = df.fillna(method='bfill')
            else:
                df = df.fillna(value)
            
            logger.debug(f"填充所有列的空值: {method or value}")
        
        return df
    
    def _remove_rows(self, df: pd.DataFrame, rule: Dict[str, Any]) -> pd.DataFrame:
        """移除行
        
        Args:
            df: 数据
            rule: 规则配置
            
        Returns:
            处理后的数据
        """
        condition = rule.get('condition')  # 条件表达式
        columns = rule.get('columns')  # 要检查的列
        values = rule.get('values')  # 要移除的值
        
        if condition:
            # 使用条件表达式
            try:
                before_count = len(df)
                df_filtered = df.query(condition)
                after_count = len(df_filtered)
                logger.info(f"条件过滤: {condition}, 移除了 {before_count - after_count} 行")
                return df_filtered
            except Exception as e:
                logger.error(f"条件表达式执行失败: {condition}, 错误: {e}")
                return df
        
        if columns and values:
            # 基于列值过滤
            if isinstance(columns, str):
                columns = [columns]
            if not isinstance(values, list):
                values = [values]
            
            df_filtered = df.copy()
            for col in columns:
                if col in df_filtered.columns:
                    mask = ~df_filtered[col].isin(values)
                    df_filtered = df_filtered[mask]
                    logger.debug(f"列 {col} 过滤值 {values}")
            
            return df_filtered
        
        logger.warning("移除行规则缺少条件或列配置")
        return df
    
    def _convert_type(self, df: pd.DataFrame, rule: Dict[str, Any]) -> pd.DataFrame:
        """类型转换
        
        Args:
            df: 数据
            rule: 规则配置
            
        Returns:
            处理后的数据
        """
        columns = rule.get('columns')
        target_type = rule.get('target_type', 'float')
        errors = rule.get('errors', 'coerce')  # 错误处理方式
        
        if isinstance(columns, str):
            columns = [columns]
        
        for col in columns:
            if col in df.columns:
                try:
                    if target_type == 'int':
                        df[col] = pd.to_numeric(df[col], errors=errors).astype('Int64')
                    elif target_type == 'float':
                        df[col] = pd.to_numeric(df[col], errors=errors)
                    elif target_type == 'datetime':
                        df[col] = pd.to_datetime(df[col], errors=errors)
                    elif target_type == 'string':
                        df[col] = df[col].astype(str)
                    else:
                        df[col] = df[col].astype(target_type)
                    
                    logger.debug(f"列 {col} 类型转换: {target_type}")
                except Exception as e:
                    logger.error(f"列 {col} 类型转换失败: {e}")
        
        return df
    
    def _strip_whitespace(self, df: pd.DataFrame, rule: Dict[str, Any]) -> pd.DataFrame:
        """去除空白字符
        
        Args:
            df: 数据
            rule: 规则配置
            
        Returns:
            处理后的数据
        """
        columns = rule.get('columns')
        
        if columns:
            if isinstance(columns, str):
                columns = [columns]
            
            for col in columns:
                if col in df.columns and df[col].dtype == 'object':
                    df[col] = df[col].str.strip()
                    logger.debug(f"列 {col} 去除空白字符")
        else:
            # 处理所有文本列
            text_columns = df.select_dtypes(include=['object']).columns
            for col in text_columns:
                df[col] = df[col].str.strip()
            
            logger.debug(f"所有文本列去除空白字符: {list(text_columns)}")
        
        return df
    
    def _normalize_text(self, df: pd.DataFrame, rule: Dict[str, Any]) -> pd.DataFrame:
        """文本标准化
        
        Args:
            df: 数据
            rule: 规则配置
            
        Returns:
            处理后的数据
        """
        columns = rule.get('columns')
        lowercase = rule.get('lowercase', False)
        uppercase = rule.get('uppercase', False)
        replace_map = rule.get('replace_map', {})  # 替换映射
        
        if isinstance(columns, str):
            columns = [columns]
        
        target_columns = columns or df.select_dtypes(include=['object']).columns
        
        for col in target_columns:
            if col in df.columns and df[col].dtype == 'object':
                if lowercase:
                    df[col] = df[col].str.lower()
                elif uppercase:
                    df[col] = df[col].str.upper()
                
                # 应用替换映射
                for old, new in replace_map.items():
                    df[col] = df[col].str.replace(old, new)
                
                logger.debug(f"列 {col} 文本标准化完成")
        
        return df
    
    def _validate_data(self, df: pd.DataFrame, rule: Dict[str, Any]) -> pd.DataFrame:
        """数据验证
        
        Args:
            df: 数据
            rule: 规则配置
            
        Returns:
            处理后的数据
        """
        columns = rule.get('columns')
        min_value = rule.get('min_value')
        max_value = rule.get('max_value')
        required = rule.get('required', False)
        
        if isinstance(columns, str):
            columns = [columns]
        
        validation_results = []
        
        for col in columns:
            if col in df.columns:
                # 检查必需值
                if required:
                    null_count = df[col].isnull().sum()
                    if null_count > 0:
                        validation_results.append(f"{col}: {null_count} 个空值")
                
                # 检查数值范围
                if min_value is not None or max_value is not None:
                    if pd.api.types.is_numeric_dtype(df[col]):
                        invalid_mask = pd.Series(False, index=df.index)
                        if min_value is not None:
                            invalid_mask |= df[col] < min_value
                        if max_value is not None:
                            invalid_mask |= df[col] > max_value
                        
                        invalid_count = invalid_mask.sum()
                        if invalid_count > 0:
                            validation_results.append(f"{col}: {invalid_count} 个值超出范围")
        
        if validation_results:
            logger.warning(f"数据验证发现问题: {', '.join(validation_results)}")
        else:
            logger.debug("数据验证通过")
        
        return df
    
    # 便捷方法
    def remove_duplicates(self, subset: Optional[List[str]] = None, keep: str = 'first'):
        """移除重复项"""
        self.add_rule('remove_duplicates', subset=subset, keep=keep)
        return self
    
    def fill_na(self, columns: Optional[Union[str, List[str]]] = None, 
                value: Any = 0, method: Optional[str] = None):
        """填充空值"""
        self.add_rule('fill_na', columns=columns, value=value, method=method)
        return self
    
    def remove_rows(self, condition: Optional[str] = None,
                   columns: Optional[Union[str, List[str]]] = None,
                   values: Optional[Any] = None):
        """移除行"""
        self.add_rule('remove_rows', condition=condition, columns=columns, values=values)
        return self
    
    def convert_type(self, columns: Union[str, List[str]], target_type: str, errors: str = 'coerce'):
        """类型转换"""
        self.add_rule('convert_type', columns=columns, target_type=target_type, errors=errors)
        return self
    
    def strip_whitespace(self, columns: Optional[Union[str, List[str]]] = None):
        """去除空白字符"""
        self.add_rule('strip_whitespace', columns=columns)
        return self
    
    def normalize_text(self, columns: Optional[Union[str, List[str]]] = None,
                      lowercase: bool = False, uppercase: bool = False,
                      replace_map: Optional[Dict[str, str]] = None):
        """文本标准化"""
        self.add_rule('normalize_text', columns=columns, lowercase=lowercase, 
                     uppercase=uppercase, replace_map=replace_map or {})
        return self
    
    def validate_data(self, columns: Union[str, List[str]], 
                     min_value: Optional[float] = None,
                     max_value: Optional[float] = None,
                     required: bool = False):
        """数据验证"""
        self.add_rule('validate_data', columns=columns, min_value=min_value,
                     max_value=max_value, required=required)
        return self