"""
HiveAgent main class - the primary user interface for pypabhiveagent.
"""

import logging
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Dict, Optional
import pandas as pd

from pypabhiveagent.models import Config, QueryResult
from pypabhiveagent.exceptions import ConfigurationError, HiveQueryError
from pypabhiveagent.hive_client import HiveClient
from pypabhiveagent.aigc_client import AIGCClient
from pypabhiveagent.result_parser import ResultParser
from pypabhiveagent.storage_manager import StorageManager

logger = logging.getLogger(__name__)


class HiveAgent:
    """
    Main facade class for querying Hive data and processing with AIGC applications.
    """
    
    def __init__(self, hive_backend='auto', **hive_kwargs):
        """
        Initialize HiveAgent instance.
        
        Args:
            hive_backend (str): Hive backend to use ('spark', 'pyhive', or 'auto')
            **hive_kwargs: Additional Hive connection parameters
        """
        self._config = {}
        self._hive_client = None
        self._aigc_client = None
        self._result_parser = None
        self._hive_backend = hive_backend
        self._hive_kwargs = hive_kwargs
        self._storage_manager = None
        logger.debug("HiveAgent instance initialized")
    
    def set_config(self, **kwargs) -> 'HiveAgent':
        """
        Configure system parameters (supports method chaining).
        
        Args:
            url: AIGC service URL
            app_id: Application ID
            token: Authentication token
            aigc_app_id: AIGC application ID
            timeout: AIGC request timeout in seconds (default: 300)
            max_retries: Maximum retries for AIGC requests (default: 3)
            sleep: Sleep time before each AIGC request in seconds (default: 2)
                   Applied before every request to avoid rate limiting
            
        Returns:
            self: Supports method chaining
        """
        self._config.update(kwargs)
        logger.debug(f"Configuration updated with {len(kwargs)} parameters")
        return self
    
    def _validate_config(self):
        """
        Validate required configuration parameters.
        
        Raises:
            ConfigurationError: If required configuration is missing
        """
        # Create Config object from _config dict
        config = Config(
            url=self._config.get('url'),
            app_id=self._config.get('app_id'),
            token=self._config.get('token'),
            aigc_app_id=self._config.get('aigc_app_id')
        )
        
        # Validate and get missing fields
        missing_fields = config.validate()
        
        if missing_fields:
            error_msg = f"Missing required configuration: {', '.join(missing_fields)}"
            logger.error(error_msg)
            raise ConfigurationError(error_msg)
        
        logger.debug("Configuration validation passed")
    
    def _validate_app_type(self, app_type: str):
        """
        Validate application type parameter.
        
        Args:
            app_type: Application type to validate
            
        Raises:
            ConfigurationError: If app_type is not 'prompt' or 'agent'
        """
        if app_type not in ['prompt', 'agent']:
            error_msg = f"Application type must be 'prompt' or 'agent', got: {app_type}"
            logger.error(error_msg)
            raise ConfigurationError(error_msg)
        
        logger.debug(f"Application type validated: {app_type}")
    
    def _normalize_field_mapping(self, field_mapping: dict, df_columns: list) -> dict:
        """
        Normalize field mapping to match DataFrame column names (case-insensitive).
        
        Args:
            field_mapping: Original field mapping dictionary with potentially mismatched case
            df_columns: List of actual DataFrame column names
            
        Returns:
            dict: Normalized field mapping with correct column names
            
        Raises:
            ConfigurationError: If mapped fields don't exist in DataFrame (case-insensitive)
        """
        if not field_mapping:
            error_msg = "field_mapping parameter is required"
            logger.error(error_msg)
            raise ConfigurationError(error_msg)
        
        # Create case-insensitive lookup for DataFrame columns
        df_columns_lower = {col.lower(): col for col in df_columns}
        
        # Normalize field mapping
        normalized_mapping = {}
        missing_fields = []
        
        for field_name, param_name in field_mapping.items():
            field_name_lower = field_name.lower()
            
            if field_name_lower in df_columns_lower:
                # Use the actual column name from DataFrame
                actual_col_name = df_columns_lower[field_name_lower]
                normalized_mapping[actual_col_name] = param_name
                logger.debug(f"Normalized field mapping: '{field_name}' -> '{actual_col_name}'")
            else:
                missing_fields.append(field_name)
        
        if missing_fields:
            error_msg = f"Mapped fields not found in data source (case-insensitive): {', '.join(missing_fields)}"
            logger.error(error_msg)
            logger.debug(f"Available columns: {', '.join(df_columns)}")
            raise ConfigurationError(error_msg)
        
        logger.info(f"Field mapping normalized: {len(normalized_mapping)} fields")
        return normalized_mapping
    
    def _validate_field_mapping(self, field_mapping: dict, df_columns: list):
        """
        Validate field mapping against DataFrame columns.
        
        Args:
            field_mapping: Field mapping dictionary (should already be normalized)
            df_columns: List of DataFrame column names
            
        Raises:
            ConfigurationError: If mapped fields don't exist in DataFrame
        """
        if not field_mapping:
            error_msg = "field_mapping parameter is required"
            logger.error(error_msg)
            raise ConfigurationError(error_msg)
        
        # Check if all mapped fields exist in DataFrame
        missing_fields = []
        for field_name in field_mapping.keys():
            if field_name not in df_columns:
                missing_fields.append(field_name)
        
        if missing_fields:
            error_msg = f"Mapped fields not found in query results: {', '.join(missing_fields)}"
            logger.error(error_msg)
            raise ConfigurationError(error_msg)
        
        logger.debug(f"Field mapping validated: {len(field_mapping)} fields")
    
    def _initialize_clients(self):
        """Initialize all client components."""
        # Initialize HiveClient (uses Spark's getOrCreate())
        self._hive_client = HiveClient()
        logger.debug("HiveClient initialized")
        
        # Initialize AIGCClient
        self._aigc_client = AIGCClient(
            url=self._config.get('url'),
            app_id=self._config.get('app_id'),
            token=self._config.get('token'),
            aigc_app_id=self._config.get('aigc_app_id'),
            timeout=self._config.get('timeout', 300),
            max_retries=self._config.get('max_retries', 3),
            sleep=self._config.get('sleep', 2)
        )
        logger.debug("AIGCClient initialized")
        
        # Initialize ResultParser
        self._result_parser = ResultParser()
        logger.debug("ResultParser initialized")
        
        # Initialize StorageManager
        self._storage_manager = StorageManager(
            aigc_app_id=self._config.get('aigc_app_id')
        )
        logger.debug("StorageManager initialized")
    
    def _build_aigc_params(self, row: pd.Series, field_mapping: dict, app_type: str) -> dict:
        """
        Build AIGC request parameters from a data row.
        
        Args:
            row: Data row
            field_mapping: Field mapping dictionary
            app_type: Application type ('prompt' or 'agent')
            
        Returns:
            dict: AIGC request parameters
        """
        # Determine if single or multiple parameters
        is_single_param = len(field_mapping) == 1
        
        if is_single_param:
            # Single parameter input
            field_name = list(field_mapping.keys())[0]
            value = row[field_name]
            
            if app_type == 'prompt':
                return {'query': str(value)}
            else:  # agent
                # For agent single parameter, use 'content' key
                return {'content': str(value)}
        else:
            # Multiple parameters input
            params_dict = {}
            for field_name, param_name in field_mapping.items():
                # Use param_name if provided, otherwise use field_name
                key = param_name if param_name else field_name
                params_dict[key] = str(row[field_name])
            
            if app_type == 'prompt':
                return {'dynamicColumMap': params_dict}
            else:  # agent
                # For agent multiple parameters, the field mapped to 'content' goes to message
                # and others go as separate parameters
                if 'content' not in params_dict:
                    # This should not happen due to validation, but handle gracefully
                    logger.error("No field mapped to 'content' for Agent multi-parameter request")
                    return {}
                
                # Build result with content and other parameters
                result = {'content': params_dict['content']}
                
                # Add other parameters (excluding 'content')
                for key, value in params_dict.items():
                    if key != 'content':
                        result[key] = value
                
                return result
    
    def _process_single_row(
        self,
        row: pd.Series,
        row_index: int,
        app_type: str,
        field_mapping: dict,
        include_reasoning: bool
    ) -> tuple:
        """
        Process a single row with AIGC application.
        
        Args:
            row: Data row to process
            row_index: Row index for logging
            app_type: Application type
            field_mapping: Field mapping
            include_reasoning: Whether to include reasoning
            
        Returns:
            tuple: (processed_row, error_info)
                - processed_row: pd.Series with results merged
                - error_info: dict with error details if failed, None if successful
        """
        try:
            # Build AIGC request parameters
            params = self._build_aigc_params(row, field_mapping, app_type)
            
            # Call AIGC application
            if app_type == 'prompt':
                response = self._aigc_client.call_prompt_app(params)
                parsed_result = self._result_parser.parse_prompt_response(response, include_reasoning)
            else:  # agent
                response = self._aigc_client.call_agent_app(params)
                parsed_result = self._result_parser.parse_agent_response(response, include_reasoning)
            
            # Merge result to row
            merged_row = self._result_parser.merge_result_to_row(row, parsed_result)
            
            logger.debug(f"Successfully processed row {row_index}")
            return merged_row, None
            
        except Exception as e:
            logger.error(f"Error processing row {row_index}: {str(e)}")
            # Create error result with consistent schema
            error_result = {
                'success': False,
                'result': None,
                'reasoning': None,
                'judge_time': '',  # Empty string for failed requests
                'log_id': '',      # Empty string for failed requests
                'error': str(e)
            }
            # Merge error result to row to maintain consistent schema
            merged_row = self._result_parser.merge_result_to_row(row, error_result)
            
            error_info = {
                'row_index': row_index,
                'error': str(e),
                'error_type': type(e).__name__
            }
            return merged_row, error_info
    
    def _process_dataframe(
        self,
        df: pd.DataFrame,
        app_type: str,
        field_mapping: dict,
        include_reasoning: bool,
        max_workers: int
    ) -> tuple:
        """
        Process all rows in DataFrame with concurrent execution.
        
        Args:
            df: DataFrame to process
            app_type: Application type
            field_mapping: Field mapping
            include_reasoning: Whether to include reasoning
            max_workers: Maximum concurrent workers
            
        Returns:
            tuple: (processed_df, errors_list, result_columns)
                - processed_df: DataFrame with all rows (failed rows have None in result columns)
                - errors_list: List of error dictionaries
                - result_columns: Set of result column names added by AIGC
        """
        logger.info(f"Starting batch processing of {len(df)} rows with {max_workers} workers")
        
        processed_rows = []
        errors = []
        result_columns = set()  # Track columns added by AIGC results
        
        # Use ThreadPoolExecutor for concurrent processing
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            # Submit all tasks
            future_to_index = {}
            for idx, row in df.iterrows():
                future = executor.submit(
                    self._process_single_row,
                    row,
                    idx,
                    app_type,
                    field_mapping,
                    include_reasoning
                )
                future_to_index[future] = idx
            
            # Collect results as they complete
            completed = 0
            for future in as_completed(future_to_index):
                idx = future_to_index[future]
                try:
                    result_row, error_info = future.result()
                    
                    # Track result columns (columns not in original data)
                    original_columns = set(df.columns)
                    new_columns = set(result_row.index) - original_columns
                    result_columns.update(new_columns)
                    
                    processed_rows.append((idx, result_row, error_info))
                    
                    # Record error if present
                    if error_info:
                        errors.append(error_info)
                    
                    completed += 1
                    if completed % 10 == 0 or completed == len(df):
                        logger.info(f"Progress: {completed}/{len(df)} rows processed")
                        
                except Exception as e:
                    logger.error(f"Failed to get result for row {idx}: {str(e)}")
                    # For unexpected errors, keep original row
                    original_row = df.loc[idx].copy()
                    error_info = {
                        'row_index': idx,
                        'error': str(e),
                        'error_type': 'UnexpectedError'
                    }
                    processed_rows.append((idx, original_row, error_info))
                    errors.append(error_info)
        
        # Sort by original index
        processed_rows.sort(key=lambda x: x[0])
        
        # Create DataFrame with unified structure
        result_df = self._unify_dataframe_structure(processed_rows, result_columns)
        
        logger.info(f"Batch processing completed: {len(result_df)} rows processed, {len(errors)} errors")
        
        return result_df, errors, result_columns
    
    def _unify_dataframe_structure(self, processed_rows: list, result_columns: set) -> pd.DataFrame:
        """
        Unify DataFrame structure by ensuring all rows have the same columns.
        Failed rows will have None values in result columns.
        
        Args:
            processed_rows: List of (index, row, error_info) tuples
            result_columns: Set of result column names
            
        Returns:
            pd.DataFrame: Unified DataFrame with consistent structure
        """
        unified_rows = []
        
        for idx, row, error_info in processed_rows:
            unified_row = row.copy()
            
            # If this row failed, ensure it has all result columns with None values
            if error_info:
                for col in result_columns:
                    if col not in unified_row.index:
                        unified_row[col] = None
            else:
                # If this row succeeded, ensure it has all result columns
                # (some might be missing if different rows return different structures)
                for col in result_columns:
                    if col not in unified_row.index:
                        unified_row[col] = None
            
            unified_rows.append(unified_row)
        
        # Create DataFrame
        result_df = pd.DataFrame(unified_rows)
        
        logger.debug(f"Unified DataFrame structure: {len(result_df)} rows, {len(result_df.columns)} columns")
        
        return result_df
    
    def query(
        self,
        sql: Optional[str] = None,
        df: Optional[pd.DataFrame] = None,
        app_type: str = None,
        field_mapping: dict = None,
        include_reasoning: bool = False,
        cache_to_excel: bool = True,
        save_to_hive: bool = True,
        result_table_name: Optional[str] = None,
        write_mode: str = 'overwrite',
        max_workers: int = 1
    ) -> dict:
        """
        Execute complete query and processing workflow.
        
        Args:
            sql: Hive SQL query statement (optional if df is provided)
            df: Pandas DataFrame as data source (optional if sql is provided)
                - Column names will be matched with field_mapping (case-insensitive)
                - Example: df with columns ['ID', 'Question', 'DT']
            app_type: Application type ('prompt' or 'agent')
            field_mapping: Field mapping dictionary
                - Keys should match DataFrame column names (case-insensitive)
                - Single parameter: {'field_name': 'param_name'} or {'field_name': None}
                - Multiple parameters: {'field1': 'param1', 'field2': 'param2'}
                - Agent multi-param: One field must map to 'content' for message content
                  Example: {'customer_name': 'content', 'order_status': 'status', 'order_amount': 'amount'}
            include_reasoning: Whether to include reasoning chain results
            cache_to_excel: Whether to cache to Excel
            save_to_hive: Whether to save to Hive
            result_table_name: Result table name (required when save_to_hive=True)
            write_mode: Write mode ('append' or 'overwrite')
            max_workers: Maximum concurrent worker threads (default: 1)
            
        Returns:
            dict: Processing result dictionary containing:
                - success: bool
                - message: str
                - rows_processed: int
                - df: pandas DataFrame with results
                - excel_path: str (if cache_to_excel=True)
                - hive_table: str (if save_to_hive=True)
                - errors: list of error dictionaries
        """
        logger.info("=" * 80)
        logger.info("Starting HiveAgent query workflow")
        logger.info("=" * 80)
        
        try:
            # Step 1: Validate configuration
            logger.info("Step 1: Validating configuration")
            self._validate_config()
            
            # Step 2: Validate application type
            logger.info("Step 2: Validating application type")
            if not app_type:
                error_msg = "app_type parameter is required ('prompt' or 'agent')"
                logger.error(error_msg)
                raise ConfigurationError(error_msg)
            self._validate_app_type(app_type)
            
            # Step 3: Validate field_mapping parameter
            if not field_mapping:
                error_msg = "field_mapping parameter is required"
                logger.error(error_msg)
                raise ConfigurationError(error_msg)
            
            # Step 4: Validate data source (sql or df)
            if sql is None and df is None:
                error_msg = "Either 'sql' or 'df' parameter must be provided"
                logger.error(error_msg)
                raise ConfigurationError(error_msg)
            
            if sql is not None and df is not None:
                error_msg = "Cannot provide both 'sql' and 'df' parameters, choose one"
                logger.error(error_msg)
                raise ConfigurationError(error_msg)
            
            # Step 5: Validate save_to_hive requirements
            if save_to_hive and not result_table_name:
                error_msg = "Result table name is required when save_to_hive is True"
                logger.error(error_msg)
                raise ConfigurationError(error_msg)
            
            # Step 6: Initialize clients
            logger.info("Step 3: Initializing clients")
            self._initialize_clients()
            
            # Step 7: Get data (from SQL or DataFrame)
            if sql is not None:
                # Execute Hive query
                logger.info("Step 4: Executing Hive query")
                if not sql.strip():
                    error_msg = "SQL statement cannot be empty"
                    logger.error(error_msg)
                    raise HiveQueryError(error_msg)
                source_df = self._hive_client.execute_query(sql)
            else:
                # Use provided DataFrame
                logger.info("Step 4: Using provided DataFrame as data source")
                if not isinstance(df, pd.DataFrame):
                    error_msg = "df parameter must be a pandas DataFrame"
                    logger.error(error_msg)
                    raise ConfigurationError(error_msg)
                source_df = df.copy()  # Make a copy to avoid modifying original
                logger.info(f"DataFrame loaded: {len(source_df)} rows, {len(source_df.columns)} columns")
            
            if source_df.empty:
                logger.warning("Data source is empty, no results to process")
                return QueryResult(
                    success=True,
                    message="Data source is empty, no results to process",
                    rows_processed=0,
                    df=pd.DataFrame()
                ).to_dict()
            
            # Step 8: Normalize field mapping for case-insensitive matching
            logger.info("Step 5: Normalizing field mapping for case-insensitive matching")
            normalized_field_mapping = self._normalize_field_mapping(field_mapping, source_df.columns.tolist())
            
            # Step 9: Validate field mapping
            logger.info("Step 6: Validating field mapping")
            self._validate_field_mapping(normalized_field_mapping, source_df.columns.tolist())
            
            # Step 10: Validate content mapping for Agent multi-parameter case
            if app_type == 'agent' and len(normalized_field_mapping) > 1:
                # Check if one of the field mappings maps to 'content'
                content_fields = [field for field, param in normalized_field_mapping.items() if param == 'content']
                if not content_fields:
                    error_msg = "For Agent application with multiple parameters, one field must map to 'content' for message content. Example: {'customer_name': 'content', 'order_status': 'status'}"
                    logger.error(error_msg)
                    raise ConfigurationError(error_msg)
                
                if len(content_fields) > 1:
                    error_msg = f"Only one field can map to 'content', found multiple: {content_fields}"
                    logger.error(error_msg)
                    raise ConfigurationError(error_msg)
                
                logger.debug(f"Content field validated: {content_fields[0]} -> content")
            
            # Step 11: Process data with AIGC
            logger.info("Step 7: Processing data with AIGC application")
            result_df, errors, result_columns = self._process_dataframe(
                df=source_df,
                app_type=app_type,
                field_mapping=normalized_field_mapping,
                include_reasoning=include_reasoning,
                max_workers=max_workers
            )
            
            # Step 12: Save to Excel if requested
            excel_path = None
            if cache_to_excel:
                logger.info("Step 8: Saving results to Excel")
                try:
                    excel_path = self._storage_manager.save_to_excel(result_df)
                    logger.info(f"Excel file saved: {excel_path}")
                except Exception as e:
                    logger.error(f"Failed to save Excel file: {str(e)}")
                    # Don't fail the entire operation if Excel save fails
            
            # Step 13: Save to Hive if requested
            hive_table = None
            if save_to_hive:
                logger.info("Step 9: Saving results to Hive")
                try:
                    hive_result = self._storage_manager.save_to_hive(
                        df=result_df,
                        hive_client=self._hive_client,
                        table_name=result_table_name,
                        mode=write_mode
                    )
                    hive_table = hive_result['table_name']
                    logger.info(f"Data saved to Hive table: {hive_table}")
                except Exception as e:
                    logger.error(f"Failed to save to Hive: {str(e)}")
                    # Don't fail the entire operation if Hive save fails
            
            # Step 14: Construct result with DataFrame
            logger.info("Step 10: Constructing result")
            
            success = len(errors) == 0
            if success:
                message = f"Successfully processed {len(result_df)} rows"
            else:
                message = f"Processed {len(result_df)} rows with {len(errors)} errors"
            
            result = QueryResult(
                success=success,
                message=message,
                rows_processed=len(result_df),
                df=result_df,
                excel_path=excel_path,
                hive_table=hive_table,
                errors=errors
            )
            
            logger.info("=" * 80)
            logger.info(f"Workflow completed: {message}")
            
            # Print failure summary if there are errors
            if errors:
                logger.warning("=" * 80)
                logger.warning(f"⚠️  FAILURE SUMMARY")
                logger.warning("=" * 80)
                logger.warning(f"Total failed rows: {len(errors)}")
                logger.warning(f"Failed row indices: {[e['row_index'] for e in errors[:10]]}" + 
                             (f" ... and {len(errors) - 10} more" if len(errors) > 10 else ""))
                logger.warning("")
                logger.warning(f"📌 Note: Failed rows are included in the result with:")
                logger.warning(f"   - Original data preserved")
                logger.warning(f"   - Result columns filled with None values")
                if result_columns:
                    logger.warning(f"   - Result columns: {', '.join(sorted(result_columns))}")
                logger.warning("")
                logger.warning(f"💡 Tip: You can identify failed rows by checking for None values in result columns")
                logger.warning("=" * 80)
                
                # Also print to console (not just log)
                print("\n" + "=" * 80)
                print(f"⚠️  PROCESSING COMPLETED WITH ERRORS")
                print("=" * 80)
                print(f"✓ Successfully processed: {len(result_df) - len(errors)} rows")
                print(f"✗ Failed to process: {len(errors)} rows")
                print(f"\n📌 Failed rows are included in the result with:")
                print(f"   - Original data preserved")
                print(f"   - Result columns filled with None values")
                if result_columns:
                    print(f"   - Result columns: {', '.join(sorted(result_columns))}")
                print(f"\n💡 Tip: Filter failed rows using: df[df['result_column'].isna()]")
                print("=" * 80 + "\n")
            else:
                logger.info("✓ All rows processed successfully")
            
            logger.info("=" * 80)
            
            return result.to_dict()
            
        except Exception as e:
            error_msg = f"Query workflow failed: {str(e)}"
            logger.error(error_msg)
            
            result = QueryResult(
                success=False,
                message=error_msg,
                rows_processed=0,
                df=pd.DataFrame(),
                errors=[{'error': str(e)}]
            )
            
            return result.to_dict()
        
        finally:
            # Clean up resources
            if self._hive_client:
                try:
                    self._hive_client.close()
                    logger.debug("HiveClient closed")
                except Exception as e:
                    logger.warning(f"Error closing HiveClient: {str(e)}")
    
    @staticmethod
    def demo():
        """
        打印使用示例代码，包含详细的入参和出参说明。
        帮助首次使用的用户快速上手。
        """
        demo_code = '''
# ============================================================================
# pypabhiveagent 快速上手指南
# ============================================================================
# 
# pypabhiveagent 是一个用于查询 Hive 数据并通过 AIGC 应用处理的 Python 包
# 
# 核心功能：
# 1. 从 Hive 查询数据 或 直接使用 pandas DataFrame
# 2. 发送到 AIGC 应用（Prompt 或 Agent）进行处理
# 3. 获取处理结果（DataFrame 格式）
# 4. 保存到 Excel 或 Hive 表
#
# ============================================================================

import pandas as pd
from pypabhiveagent import HiveAgent


# ============================================================================
# 第一部分：基础配置
# ============================================================================

# 创建 HiveAgent 实例
agent = HiveAgent()

# 配置 AIGC 服务连接参数（必需）
agent.set_config(
    url="https://aigc-api.example.com/prompt",    # AIGC 服务地址
    app_id="your_app_id",                          # 应用 ID
    token="your_token",                            # 认证 token
    aigc_app_id="prompt_app_12345",                # AIGC 应用 ID
    
    # 以下为可选参数（有默认值）
    timeout=300,        # 请求超时时间（秒），默认 300
    max_retries=3,     # 最大重试次数，默认 3
    sleep=2            # 请求间隔时间（秒），默认 2
)


# ============================================================================
# 第二部分：核心方法 - query()
# ============================================================================
#
# query() 方法是核心接口，用于执行完整的查询和处理流程
#
# ============================================================================

# ----------------------------------------------------------------------------
# 入参说明（Parameters）
# ----------------------------------------------------------------------------
#
# 数据源（二选一，必需）：
#   sql: str                    - Hive SQL 查询语句
#   df: pd.DataFrame            - pandas DataFrame（直接传入数据）
#
# 处理配置（必需）：
#   app_type: str               - 应用类型，'prompt' 或 'agent'
#   field_mapping: dict         - 字段映射，指定哪些列发送到 AIGC
#                                 格式：{'列名': '参数名'} 或 {'列名': None}
#
# 可选参数（有默认值）：
#   include_reasoning: bool     - 是否包含推理过程，默认 False
#   cache_to_excel: bool        - 是否保存到 Excel，默认 True
#   save_to_hive: bool          - 是否保存到 Hive，默认 True
#   result_table_name: str      - 结果表名（save_to_hive=True 时必需）
#   write_mode: str             - 写入模式，'append' 或 'overwrite'，默认 'overwrite'
#   max_workers: int            - 并发线程数，默认 1
#
# ----------------------------------------------------------------------------
# 出参说明（Returns）
# ----------------------------------------------------------------------------
#
# 返回一个字典（dict），包含以下字段：
#
#   success: bool               - 处理是否成功（True/False）
#   message: str                - 处理结果消息
#   rows_processed: int         - 处理的行数
#   df: pd.DataFrame            - 处理后的完整 DataFrame（重要！）
#   excel_path: str             - Excel 文件路径（如果 cache_to_excel=True）
#   hive_table: str             - Hive 表名（如果 save_to_hive=True）
#   errors: list                - 错误列表（如果有失败的行）
#
# 重点：result['df'] 包含完整的处理结果，可以直接使用！
#
# ============================================================================


# ============================================================================
# 示例 1：最简单的使用方式（从 Hive 查询）
# ============================================================================

result = agent.query(
    # 数据源：SQL 查询
    sql="SELECT id, question, dt FROM qa_table WHERE dt='20231201' LIMIT 10",
    
    # 处理配置
    app_type='prompt',                    # 使用 Prompt 应用
    field_mapping={'question': None},     # 将 question 列发送到 AIGC
    
    # 存储配置（可选）
    save_to_hive=True,
    result_table_name='qa_results'
)

# 查看结果
print(f"✓ 处理状态: {result['success']}")
print(f"✓ 处理消息: {result['message']}")
print(f"✓ 处理行数: {result['rows_processed']}")

# 获取结果 DataFrame（重要！）
result_df = result['df']
print(f"✓ 结果 DataFrame: {len(result_df)} 行 x {len(result_df.columns)} 列")
print(result_df.head())


# ============================================================================
# 示例 2：使用 DataFrame 作为输入（推荐）
# ============================================================================

# 准备数据（可以从 CSV、Excel、API 等加载）
input_df = pd.DataFrame({
    'id': [1, 2, 3, 4, 5],
    'question': ['问题1', '问题2', '问题3', '问题4', '问题5'],
    'dt': ['20231201'] * 5
})

result = agent.query(
    # 数据源：DataFrame（而不是 SQL）
    df=input_df,
    
    # 处理配置
    app_type='prompt',
    field_mapping={'question': None},
    
    # 存储配置
    save_to_hive=True,
    result_table_name='qa_results'
)

# 获取结果 DataFrame
result_df = result['df']
print(result_df)


# ============================================================================
# 示例 3：完整的入参和出参示例
# ============================================================================

# 准备输入数据
input_df = pd.DataFrame({
    'id': [1, 2, 3],
    'question': ['什么是AI?', '什么是ML?', '什么是DL?'],
    'category': ['技术', '技术', '技术'],
    'dt': ['20231201'] * 3
})

# 执行处理
result = agent.query(
    # ========== 入参 ==========
    # 数据源
    df=input_df,                          # 使用 DataFrame 输入
    
    # 处理配置
    app_type='prompt',                    # 应用类型
    field_mapping={'question': None},     # 字段映射
    
    # 可选配置
    include_reasoning=False,              # 不包含推理过程
    cache_to_excel=True,                  # 保存到 Excel
    save_to_hive=True,                    # 保存到 Hive
    result_table_name='qa_results',       # 结果表名
    write_mode='overwrite',                  # 追加模式
    max_workers=5                         # 5 个并发线程
)

# ========== 出参 ==========
print("\\n处理结果：")
print(f"1. success: {result['success']}")              # True/False
print(f"2. message: {result['message']}")              # 处理消息
print(f"3. rows_processed: {result['rows_processed']}")# 处理行数

# 最重要：获取结果 DataFrame
result_df = result['df']
print(f"\\n4. df (DataFrame):")
print(f"   - 行数: {len(result_df)}")
print(f"   - 列数: {len(result_df.columns)}")
print(f"   - 列名: {list(result_df.columns)}")
print(f"\\n   数据预览:")
print(result_df)

# 其他输出
if result.get('excel_path'):
    print(f"\\n5. excel_path: {result['excel_path']}")
if result.get('hive_table'):
    print(f"6. hive_table: {result['hive_table']}")
if result.get('errors'):
    print(f"7. errors: {len(result['errors'])} 个错误")


# ============================================================================
# 示例 4：多参数输入
# ============================================================================

# 准备多列数据
input_df = pd.DataFrame({
    'user_name': ['张三', '李四', '王五'],
    'product_name': ['产品A', '产品B', '产品C'],
    'amount': [100, 200, 300],
    'dt': ['20231201'] * 3
})

result = agent.query(
    df=input_df,
    app_type='prompt',
    
    # 多参数映射：将多个列发送到 AIGC
    field_mapping={
        'user_name': 'userName',        # 列名 -> AIGC 参数名
        'product_name': 'productName',
        'amount': 'amount'
    },
    
    save_to_hive=False,
    cache_to_excel=True
)

result_df = result['df']
print(result_df)


# ============================================================================
# 示例 5：处理结果 DataFrame 的常见操作
# ============================================================================

result = agent.query(
    sql="SELECT id, question FROM qa_table WHERE dt='20231201'",
    app_type='prompt',
    field_mapping={'question': None}
)

# 获取结果 DataFrame
result_df = result['df']

# 操作 1：查看数据
print("前 5 行:")
print(result_df.head())

# 操作 2：筛选数据
if 'confidence' in result_df.columns:
    high_quality = result_df[result_df['confidence'].astype(float) > 0.9]
    print(f"\\n高质量结果: {len(high_quality)} 行")

# 操作 3：保存到文件
result_df.to_csv('output.csv', index=False)
result_df.to_excel('output.xlsx', index=False)
result_df.to_json('output.json', orient='records')
print("\\n✓ 已保存到 CSV、Excel 和 JSON")

# 操作 4：数据分析
print(f"\\n数据统计:")
print(result_df.describe())


# ============================================================================
# 示例 6：错误处理
# ============================================================================

result = agent.query(
    sql="SELECT id, question FROM qa_table WHERE dt='20231201'",
    app_type='prompt',
    field_mapping={'question': None}
)

result_df = result['df']

# 检查是否有失败的行
# 失败的行在结果列中为 None
if 'answer' in result_df.columns:
    success_rows = result_df[result_df['answer'].notna()]
    failed_rows = result_df[result_df['answer'].isna()]
    
    print(f"成功: {len(success_rows)} 行")
    print(f"失败: {len(failed_rows)} 行")
    
    # 查看失败的行
    if not failed_rows.empty:
        print("\\n失败的行:")
        print(failed_rows)

# 查看错误详情
if result.get('errors'):
    print(f"\\n错误详情:")
    for error in result['errors'][:5]:
        print(f"  行 {error['row_index']}: {error['error']}")


# ============================================================================
# 示例 7：Agent 应用（多参数）
# ============================================================================

input_df = pd.DataFrame({
    'customer_name': ['张三', '李四'],
    'order_status': ['pending', 'processing'],
    'order_amount': [1000, 2000],
    'dt': ['20231201'] * 2
})

result = agent.query(
    df=input_df,
    app_type='agent',  # 使用 Agent 应用
    
    # Agent 多参数：必须有一个字段映射到 'content'
    field_mapping={
        'customer_name': 'content',    # 作为消息内容
        'order_status': 'status',      # 其他参数
        'order_amount': 'amount'
    },
    
    include_reasoning=True,  # 包含推理过程
    save_to_hive=True,
    result_table_name='order_analysis'
)

result_df = result['df']
print(result_df)


# ============================================================================
# 示例 8：从 CSV 到 Hive 的完整流程
# ============================================================================

# 步骤 1: 加载 CSV 数据
df = pd.read_csv('questions.csv')
print(f"步骤 1: 加载了 {len(df)} 行数据")

# 步骤 2: 数据预处理
df = df.dropna()                          # 删除空值
df = df.drop_duplicates()                 # 去重
df['question'] = df['question'].str.strip()  # 清理文本
print(f"步骤 2: 预处理后 {len(df)} 行")

# 步骤 3: 处理数据
result = agent.query(
    df=df,
    app_type='prompt',
    field_mapping={'question': None},
    save_to_hive=True,
    result_table_name='qa_results',
    max_workers=5
)

# 步骤 4: 查看结果
print(f"\\n步骤 4: 处理完成")
print(f"  状态: {result['success']}")
print(f"  处理行数: {result['rows_processed']}")

# 步骤 5: 分析结果
result_df = result['df']
if 'answer' in result_df.columns:
    success_rate = (result_df['answer'].notna().sum() / len(result_df)) * 100
    print(f"  成功率: {success_rate:.1f}%")


# ============================================================================
# 重要提示
# ============================================================================
#
# 1. 数据源：
#    - 使用 sql 参数：从 Hive 查询数据
#    - 使用 df 参数：直接传入 DataFrame（推荐）
#    - 两者只能选一个
#
# 2. 结果获取：
#    - result['df'] 是最重要的输出
#    - 包含原始数据 + AIGC 处理结果
#    - 可以直接用于后续分析
#
# 3. 字段映射：
#    - 列名大小写不敏感（'Question' 可以匹配 'question'）
#    - 单参数：{'question': None}
#    - 多参数：{'field1': 'param1', 'field2': 'param2'}
#
# 4. 错误处理：
#    - 失败的行保留原始数据
#    - 结果列填充 None 值
#    - 通过 result['df']['column'].isna() 识别失败行
#
# 5. 性能优化：
#    - max_workers: 控制并发数（默认 1）
#    - timeout: 控制超时时间（默认 300 秒）
#    - sleep: 控制请求间隔（默认 2 秒）
#
# 6. 默认参数：
#    - timeout=300（秒）
#    - max_retries=3（次）
#    - sleep=2（秒）
#    - max_workers=1（个）
#    - include_reasoning=False
#    - cache_to_excel=True
#    - save_to_hive=True
#    - write_mode='overwrite'
#
# ============================================================================


# ============================================================================
# 快速参考
# ============================================================================
#
# 最简单的使用：
#
#   agent = HiveAgent()
#   agent.set_config(url="...", app_id="...", token="...", aigc_app_id="...")
#   
#   result = agent.query(
#       df=my_dataframe,
#       app_type='prompt',
#       field_mapping={'text': None}
#   )
#   
#   result_df = result['df']  # 获取结果
#
# ============================================================================


'''
        print(demo_code)
