"""
StorageManager - manages Excel caching and Hive storage operations.
"""

import logging
import os
from datetime import datetime
from typing import Optional
import pandas as pd
from pypabhiveagent.exceptions import StorageError

logger = logging.getLogger(__name__)


class StorageManager:
    """
    Storage manager that handles Excel caching and Hive storage coordination.
    """
    
    def __init__(self, aigc_app_id: str):
        """
        Initialize storage manager.
        
        Args:
            aigc_app_id: AIGC application ID for file naming
        """
        self.aigc_app_id = aigc_app_id
        logger.debug(f"StorageManager initialized with aigc_app_id: {aigc_app_id}")
    
    def _generate_excel_filename(self) -> str:
        """
        Generate Excel filename in format: {aigcAppId}-{年月日时分秒}.xlsx
        
        Returns:
            str: Generated filename
        """
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        filename = f"{self.aigc_app_id}-{timestamp}.xlsx"
        logger.debug(f"Generated Excel filename: {filename}")
        return filename
    
    def save_to_excel(self, df: pd.DataFrame, output_dir: str = '.') -> str:
        """
        Save DataFrame to Excel file.
        
        Args:
            df: DataFrame to save
            output_dir: Output directory path (default: current directory)
            
        Returns:
            str: Full path to the saved Excel file
            
        Raises:
            StorageError: If Excel save operation fails
        """
        if df is None or df.empty:
            error_msg = "DataFrame is empty, cannot save to Excel"
            logger.warning(error_msg)
            raise StorageError(error_msg)
        
        try:
            # Generate filename
            filename = self._generate_excel_filename()
            
            # Ensure output directory exists
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
                logger.debug(f"Created output directory: {output_dir}")
            
            # Build full file path
            file_path = os.path.join(output_dir, filename)
            
            logger.info(f"Saving {len(df)} rows to Excel file: {file_path}")
            
            # Save to Excel using openpyxl engine
            df.to_excel(file_path, index=False, engine='openpyxl')
            
            logger.info(f"Successfully saved Excel file: {file_path}")
            return file_path
            
        except Exception as e:
            error_msg = f"Failed to save DataFrame to Excel: {str(e)}"
            logger.error(error_msg)
            raise StorageError(error_msg) from e
    
    def save_to_hive(
        self,
        df: pd.DataFrame,
        hive_client,
        table_name: str,
        mode: str = 'append'
    ) -> dict:
        """
        Save DataFrame to Hive table with automatic table creation if needed.
        
        Args:
            df: DataFrame to save
            hive_client: HiveClient instance for database operations
            table_name: Target table name
            mode: Write mode - 'append' or 'overwrite' (default: 'append')
            
        Returns:
            dict: Save result containing:
                - success: bool
                - rows_written: int
                - table_name: str
                - message: str
                
        Raises:
            StorageError: If Hive save operation fails
        """
        if df is None or df.empty:
            error_msg = "DataFrame is empty, cannot save to Hive"
            logger.warning(error_msg)
            raise StorageError(error_msg)
        
        if not table_name:
            error_msg = "Table name is required for Hive storage"
            logger.error(error_msg)
            raise StorageError(error_msg)
        
        try:
            logger.info(f"Preparing to save {len(df)} rows to Hive table: {table_name}")
            
            # Check if table exists
            table_exists = hive_client.table_exists(table_name)
            
            if not table_exists:
                logger.info(f"Table {table_name} does not exist, creating it")
                
                # Infer schema from DataFrame
                schema = {}
                for col_name in df.columns:
                    pandas_dtype = df[col_name].dtype
                    hive_type = hive_client._infer_hive_type(pandas_dtype)
                    schema[col_name] = hive_type
                
                logger.debug(f"Inferred schema: {schema}")
                
                # Detect if we should create a partitioned table (check for 'dt' column)
                partition_by = hive_client._detect_partition_field(df)
                
                # Create the table
                hive_client.create_table(
                    table_name=table_name,
                    schema=schema,
                    partition_by=partition_by
                )
                
                logger.info(f"Table {table_name} created successfully")
            else:
                logger.info(f"Table {table_name} already exists")
            
            # Write DataFrame to Hive
            rows_written = hive_client.write_dataframe(
                df=df,
                table_name=table_name,
                mode=mode
            )
            
            result = {
                'success': True,
                'rows_written': rows_written,
                'table_name': table_name,
                'message': f"Successfully wrote {rows_written} rows to {table_name}"
            }
            
            logger.info(result['message'])
            return result
            
        except Exception as e:
            error_msg = f"Failed to save DataFrame to Hive table {table_name}: {str(e)}"
            logger.error(error_msg)
            raise StorageError(error_msg) from e
