"""
HiveClient - handles all Hive database operations using Spark 3 or PyHive.
"""

import logging
import pandas as pd
from typing import Optional, Union
from pypabhiveagent.exceptions import HiveQueryError

# Try to import Spark dependencies
try:
    from pyspark.sql import SparkSession, DataFrame as SparkDataFrame, Row
    from pyspark.sql.types import StructType, StructField, StringType
    SPARK_AVAILABLE = True
except ImportError:
    SPARK_AVAILABLE = False
    SparkSession = None
    SparkDataFrame = None

# Try to import PyHive dependencies
try:
    from pyhive import hive
    from thrift.transport import TSocket
    from thrift.transport import TTransport
    from thrift.protocol import TBinaryProtocol
    PYHIVE_AVAILABLE = True
except ImportError:
    PYHIVE_AVAILABLE = False
    hive = None

logger = logging.getLogger(__name__)


class HiveClient:
    """
    Hive database client that handles all Hive-related operations.
    Supports both Spark and PyHive backends with automatic fallback.
    """
    
    def __init__(self, backend='auto', **kwargs):
        """
        Initialize Hive client with specified backend.
        
        Args:
            backend (str): Backend to use ('spark', 'pyhive', or 'auto')
            **kwargs: Additional connection parameters
        
        Raises:
            HiveQueryError: If no suitable backend is available
        """
        self._spark = None
        self._pyhive_conn = None
        self._backend = None
        self._connection_params = kwargs
        
        if backend == 'auto':
            # Try Spark first, then PyHive
            if SPARK_AVAILABLE:
                try:
                    self._init_spark()
                    self._backend = 'spark'
                except Exception as e:
                    logger.warning(f"Spark initialization failed: {e}")
                    if PYHIVE_AVAILABLE:
                        self._init_pyhive()
                        self._backend = 'pyhive'
                    else:
                        raise HiveQueryError("No Hive backend available. Install pyspark or pyhive.")
            elif PYHIVE_AVAILABLE:
                self._init_pyhive()
                self._backend = 'pyhive'
            else:
                raise HiveQueryError("No Hive backend available. Install pyspark or pyhive.")
        elif backend == 'spark':
            if not SPARK_AVAILABLE:
                raise HiveQueryError("PySpark not available. Install with: pip install pypabhiveagent[hive]")
            self._init_spark()
            self._backend = 'spark'
        elif backend == 'pyhive':
            if not PYHIVE_AVAILABLE:
                raise HiveQueryError("PyHive not available. Install with: pip install pypabhiveagent[hive]")
            self._init_pyhive()
            self._backend = 'pyhive'
        else:
            raise HiveQueryError(f"Unknown backend: {backend}")
    
    def _init_spark(self):
        """Initialize Spark session."""
        try:
            logger.info("Initializing SparkSession using getOrCreate()")
            self._initialize_spark_session()
            logger.info("SparkSession initialized successfully")
        except Exception as e:
            error_msg = f"Failed to initialize SparkSession: {str(e)}"
            logger.error(error_msg)
            raise HiveQueryError(error_msg) from e
    
    def _init_pyhive(self):
        """Initialize PyHive connection."""
        try:
            logger.info("Initializing PyHive connection")
            host = self._connection_params.get('host', 'localhost')
            port = self._connection_params.get('port', 10000)
            username = self._connection_params.get('username', 'hive')
            database = self._connection_params.get('database', 'default')
            
            self._pyhive_conn = hive.Connection(
                host=host,
                port=port,
                username=username,
                database=database
            )
            logger.info("PyHive connection initialized successfully")
        except Exception as e:
            error_msg = f"Failed to initialize PyHive connection: {str(e)}"
            logger.error(error_msg)
            raise HiveQueryError(error_msg) from e
    
    def _initialize_spark_session(self):
        """
        Initialize SparkSession using Spark 3's getOrCreate() method.
        This will reuse an existing SparkSession if available, or create a new one.
        
        Raises:
            Exception: If SparkSession creation fails
        """
        try:
            logger.debug("Getting or creating SparkSession with Hive support")
            
            # Use Spark 3's getOrCreate() - it will reuse existing session or create new one
            self._spark = SparkSession.builder \
                .appName("pypabhiveagent") \
                .enableHiveSupport() \
                .getOrCreate()
            
            logger.info("SparkSession obtained successfully")
            logger.debug(f"Spark version: {self._spark.version}")
                
        except Exception as e:
            logger.error(f"Failed to get or create SparkSession: {str(e)}")
            raise
    
    def execute_query(self, sql: str) -> pd.DataFrame:
        """
        Execute SQL query and return results as pandas DataFrame.
        
        Args:
            sql: SQL query statement
            
        Returns:
            pd.DataFrame: Query results
            
        Raises:
            HiveQueryError: If query execution fails or SQL is empty
        """
        if not sql or not sql.strip():
            error_msg = "SQL statement is required"
            logger.error(error_msg)
            raise HiveQueryError(error_msg)
        
        try:
            logger.info(f"Executing Hive query using {self._backend}: {sql[:100]}...")
            
            if self._backend == 'spark':
                return self._execute_query_spark(sql)
            elif self._backend == 'pyhive':
                return self._execute_query_pyhive(sql)
            else:
                raise HiveQueryError(f"Unknown backend: {self._backend}")
            
        except Exception as e:
            error_msg = f"Hive query execution failed: {str(e)}"
            logger.error(error_msg)
            raise HiveQueryError(error_msg) from e
    
    def _execute_query_spark(self, sql: str) -> pd.DataFrame:
        """Execute query using Spark backend."""
        # Execute query using spark.sql()
        spark_df = self._spark.sql(sql)
        
        # Convert to pandas DataFrame
        pandas_df = spark_df.toPandas()
        
        logger.info(f"Query executed successfully, returned {len(pandas_df)} rows")
        logger.debug(f"Query result columns: {list(pandas_df.columns)}")
        
        return pandas_df
    
    def _execute_query_pyhive(self, sql: str) -> pd.DataFrame:
        """Execute query using PyHive backend."""
        cursor = self._pyhive_conn.cursor()
        try:
            cursor.execute(sql)
            
            # Fetch results
            results = cursor.fetchall()
            columns = [desc[0] for desc in cursor.description] if cursor.description else []
            
            # Convert to pandas DataFrame
            pandas_df = pd.DataFrame(results, columns=columns)
            
            logger.info(f"Query executed successfully, returned {len(pandas_df)} rows")
            logger.debug(f"Query result columns: {list(pandas_df.columns)}")
            
            return pandas_df
        finally:
            cursor.close()
    
    def table_exists(self, table_name: str) -> bool:
        """
        Check if a table exists in the current database using multiple compatible methods.
        
        Args:
            table_name: Name of the table to check
            
        Returns:
            bool: True if table exists, False otherwise
            
        Raises:
            HiveQueryError: If the check operation fails
        """
        try:
            logger.debug(f"Checking if table exists: {table_name}")
            
            if self._backend == 'spark':
                return self._table_exists_spark(table_name)
            elif self._backend == 'pyhive':
                return self._table_exists_pyhive(table_name)
            else:
                raise HiveQueryError(f"Unknown backend: {self._backend}")
            
        except Exception as e:
            error_msg = f"Failed to check if table exists: {str(e)}"
            logger.error(error_msg)
            raise HiveQueryError(error_msg) from e
    
    def _table_exists_spark(self, table_name: str) -> bool:
        """Check table existence using Spark backend."""
        # Method 1: Try using catalog API (most reliable)
        try:
            self._spark.catalog.getTable(table_name)
            logger.debug(f"Table {table_name} exists (confirmed by catalog)")
            return True
        except Exception:
            logger.debug(f"Table {table_name} not found in catalog, trying alternative methods")
        
        # Method 2: Try DESCRIBE TABLE (works even if SHOW TABLES has issues)
        try:
            self._spark.sql(f"DESCRIBE TABLE {table_name}")
            logger.debug(f"Table {table_name} exists (confirmed by DESCRIBE)")
            return True
        except Exception:
            logger.debug(f"DESCRIBE failed for {table_name}, trying SHOW TABLES")
        
        # Method 3: Use SHOW TABLES as fallback
        try:
            result = self._spark.sql(f"SHOW TABLES")
            tables_df = result.toPandas()
            # Check if table exists in the results (case-insensitive)
            table_exists = any(
                row['tableName'].lower() == table_name.lower() 
                for _, row in tables_df.iterrows()
            )
            logger.debug(f"Table {table_name} exists: {table_exists} (confirmed by SHOW TABLES)")
            return table_exists
        except Exception as e:
            logger.warning(f"All table existence check methods failed: {str(e)}")
            return False
    
    def _table_exists_pyhive(self, table_name: str) -> bool:
        """Check table existence using PyHive backend."""
        try:
            # Try DESCRIBE TABLE first
            self.execute_query(f"DESCRIBE {table_name}")
            logger.debug(f"Table {table_name} exists (confirmed by DESCRIBE)")
            return True
        except Exception:
            try:
                # Fallback to SHOW TABLES
                result = self.execute_query("SHOW TABLES")
                table_exists = any(
                    row.lower() == table_name.lower() 
                    for row in result.iloc[:, 0] if isinstance(row, str)
                )
                logger.debug(f"Table {table_name} exists: {table_exists} (confirmed by SHOW TABLES)")
                return table_exists
            except Exception as e:
                logger.warning(f"Table existence check failed: {str(e)}")
                return False
    
    def get_table_schema(self, table_name: str) -> list:
        """
        Get the schema (column definitions) of a table.
        
        Args:
            table_name: Name of the table
            
        Returns:
            list: List of tuples (column_name, data_type, comment)
            
        Raises:
            HiveQueryError: If schema retrieval fails
        """
        try:
            logger.debug(f"Getting schema for table: {table_name}")
            
            # Use DESCRIBE to get table schema
            result = self._spark.sql(f"DESCRIBE {table_name}")
            schema_df = result.toPandas()
            
            # Extract column information (filter out partition info and empty rows)
            schema = []
            for _, row in schema_df.iterrows():
                col_name = row['col_name'].strip()
                # Stop at partition information or empty lines
                if col_name.startswith('#') or col_name == '' or col_name.lower() == 'partition information':
                    break
                data_type = row['data_type'].strip() if 'data_type' in row else ''
                comment = row['comment'].strip() if 'comment' in row and pd.notna(row['comment']) else ''
                schema.append((col_name, data_type, comment))
            
            logger.debug(f"Retrieved schema for {table_name}: {len(schema)} columns")
            return schema
            
        except Exception as e:
            error_msg = f"Failed to get table schema: {str(e)}"
            logger.error(error_msg)
            raise HiveQueryError(error_msg) from e
    
    def create_table(
        self,
        table_name: str,
        schema: dict,
        partition_by: Optional[str] = None
    ):
        """
        Create a new Hive table with the specified schema.
        All fields are explicitly created as string type for compatibility.
        
        Args:
            table_name: Name of the table to create
            schema: Dictionary mapping column names to data types (will be converted to string)
                   Example: {'id': 'int', 'name': 'string', 'dt': 'string'}
            partition_by: Optional partition column name (typically 'dt' or 'DT')
            
        Raises:
            HiveQueryError: If table creation fails
        """
        try:
            logger.info(f"Creating table: {table_name}")
            
            # Build column definitions (exclude partition column from regular columns)
            # All fields are explicitly created as string type for compatibility
            columns = []
            for col_name, _ in schema.items():
                if partition_by and col_name.lower() == partition_by.lower():
                    continue  # Skip partition column in regular columns
                # Explicitly specify string type for all columns
                columns.append(f"`{col_name}` string")
            
            columns_str = ", ".join(columns)
            
            # Build CREATE TABLE statement
            create_sql = f"CREATE TABLE IF NOT EXISTS {table_name} ({columns_str})"
            
            # Add partition clause if specified (also explicitly as string)
            if partition_by:
                create_sql += f" PARTITIONED BY (`{partition_by}` string)"
                logger.debug(f"Creating partitioned table with partition column: {partition_by} (string type)")
            
            # Add default storage format
            create_sql += " STORED AS PARQUET"
            
            logger.info(f"Creating table with all fields as string type")
            logger.debug(f"Executing CREATE TABLE: {create_sql}")
            self._spark.sql(create_sql)
            
            logger.info(f"Table {table_name} created successfully with all string fields")
            
        except Exception as e:
            error_msg = f"Failed to create table {table_name}: {str(e)}"
            logger.error(error_msg)
            raise HiveQueryError(error_msg) from e
    
    def _detect_partition_field(self, df: pd.DataFrame) -> Optional[str]:
        """
        Detect if the DataFrame has a 'dt' or 'DT' field for partitioning.
        
        Args:
            df: DataFrame to check
            
        Returns:
            Optional[str]: 'dt' or 'DT' if the field exists, None otherwise
        """
        # Check for both 'dt' and 'DT' (case-insensitive)
        for col in df.columns:
            if col.lower() == 'dt':
                logger.debug(f"Detected '{col}' field for partitioning")
                return col
        logger.debug("No 'dt' or 'DT' field found, table will not be partitioned")
        return None
    
    def _get_table_column_order(self, table_name: str) -> list:
        """
        Get the column order of a Hive table to match DataFrame column order.
        
        Args:
            table_name: Name of the table
            
        Returns:
            list: List of column names in the order they appear in the Hive table
        """
        try:
            schema = self.get_table_schema(table_name)
            column_names = [col[0] for col in schema]  # Extract column names from schema tuples
            logger.debug(f"Table {table_name} column order: {column_names}")
            return column_names
        except Exception as e:
            logger.warning(f"Could not get table column order for {table_name}: {str(e)}")
            return []
    
    def _prepare_dataframe_for_hive(self, df: pd.DataFrame, table_name: Optional[str] = None) -> pd.DataFrame:
        """
        Prepare pandas DataFrame for Hive by:
        1. Converting all columns to string type
        2. Reordering columns to match Hive table order (if table exists) - ensuring exact match
        3. Putting partition field (dt/DT) at the end if no table order is available
        
        Args:
            df: Original pandas DataFrame
            table_name: Optional table name to match column order
            
        Returns:
            pd.DataFrame: Prepared DataFrame with string types and proper column ordering
        """
        # Create a copy to avoid modifying the original DataFrame
        prepared_df = df.copy()
        
        # Convert all columns to string type for compatibility
        for col in prepared_df.columns:
            # Handle NaN/None values properly
            prepared_df[col] = prepared_df[col].fillna('').astype(str)
            logger.debug(f"Converted column '{col}' to string type")
        
        # Try to match Hive table column order if table name is provided
        if table_name and self.table_exists(table_name):
            try:
                table_columns = self._get_table_column_order(table_name)
                if table_columns:
                    # Case-insensitive column matching
                    df_columns_lower = {col.lower(): col for col in prepared_df.columns}
                    table_columns_lower = {col.lower(): col for col in table_columns}
                    
                    # Reorder DataFrame columns to match table order exactly
                    # Only include columns that exist in both DataFrame and table (case-insensitive)
                    matching_columns = []
                    for table_col_lower, table_col in table_columns_lower.items():
                        if table_col_lower in df_columns_lower:
                            df_col = df_columns_lower[table_col_lower]
                            matching_columns.append(df_col)
                    
                    # Find extra columns in DataFrame that are not in table
                    extra_columns = [col for col in prepared_df.columns 
                                   if col.lower() not in table_columns_lower]
                    
                    if matching_columns:
                        # Reorder to match table exactly, then add any extra columns
                        final_column_order = matching_columns + extra_columns
                        prepared_df = prepared_df[final_column_order]
                        logger.info(f"Reordered DataFrame columns to match table {table_name}")
                        logger.debug(f"Column order: {final_column_order}")
                    else:
                        logger.warning(f"No matching columns found between DataFrame and table {table_name}")
            except Exception as e:
                logger.warning(f"Could not match table column order: {str(e)}, using default ordering")
        
        # If no table order matching was done, use default partition field ordering
        if table_name is None or not self.table_exists(table_name):
            partition_field = self._detect_partition_field(prepared_df)
            if partition_field:
                # Reorder columns: put partition field at the end
                other_columns = [col for col in prepared_df.columns if col != partition_field]
                reordered_columns = other_columns + [partition_field]
                prepared_df = prepared_df[reordered_columns]
                logger.debug(f"Reordered columns with partition field '{partition_field}' at the end")
        
        logger.info(f"Prepared DataFrame: {len(prepared_df)} rows, {len(prepared_df.columns)} columns (all string type)")
        return prepared_df
    
    def _pandas_to_spark_with_row(self, df: pd.DataFrame) -> SparkDataFrame:
        """
        Convert pandas DataFrame to Spark DataFrame using explicit schema to avoid type inference issues.
        
        Args:
            df: Pandas DataFrame (should already be prepared with string types)
            
        Returns:
            SparkDataFrame: Converted Spark DataFrame
        """
        if df.empty:
            # Handle empty DataFrame case with explicit schema
            columns = list(df.columns) if not df.empty else []
            schema = StructType([StructField(col, StringType(), True) for col in columns])
            return self._spark.createDataFrame([], schema=schema)
        
        # Create explicit schema - all fields as StringType to avoid inference issues
        columns = list(df.columns)
        schema = StructType([StructField(col, StringType(), True) for col in columns])
        
        # Convert pandas DataFrame to list of Row objects
        rows = []
        for _, row in df.iterrows():
            # Ensure all values are strings and handle None/NaN values
            row_values = []
            for col in columns:
                value = row[col]
                if pd.isna(value) or value is None:
                    row_values.append(None)
                else:
                    row_values.append(str(value))
            rows.append(Row(*row_values))
        
        # Create Spark DataFrame with explicit schema
        spark_df = self._spark.createDataFrame(rows, schema=schema)
        logger.debug(f"Converted pandas DataFrame to Spark DataFrame with explicit string schema")
        
        return spark_df
    
    def _infer_hive_type(self, pandas_dtype) -> str:
        """
        Infer Hive data type from pandas dtype.
        For compatibility, all fields are created as string type.
        
        Args:
            pandas_dtype: Pandas data type
            
        Returns:
            str: Always returns 'string' for compatibility
        """
        # For compatibility, all fields are created as string type
        return 'string'
    
    def write_dataframe(
        self,
        df: pd.DataFrame,
        table_name: str,
        partition_value: Optional[str] = None,
        mode: str = 'append'
    ) -> int:
        """
        Write DataFrame to Hive table with support for partitioned tables.
        DataFrame will be prepared with string types and proper column ordering.
        
        Args:
            df: DataFrame to write
            table_name: Target table name
            partition_value: Partition value (for partitioned tables, typically dt value)
                           If None, will be extracted from df['dt'] or df['DT'] if it exists
            mode: Write mode - 'append' or 'overwrite' (default: 'append')
            
        Returns:
            int: Number of rows written
            
        Raises:
            HiveQueryError: If write operation fails
        """
        if df is None or df.empty:
            logger.warning("DataFrame is empty, nothing to write")
            return 0
        
        if mode not in ['append', 'overwrite']:
            error_msg = f"Invalid write mode: {mode}. Must be 'append' or 'overwrite'"
            logger.error(error_msg)
            raise HiveQueryError(error_msg)
        
        try:
            row_count = len(df)
            logger.info(f"Writing {row_count} rows to table {table_name} in {mode} mode")
            
            # Prepare DataFrame: convert to string types and reorder columns to match table
            prepared_df = self._prepare_dataframe_for_hive(df, table_name)
            
            # Convert pandas DataFrame to Spark DataFrame using Row objects for compatibility
            spark_df = self._pandas_to_spark_with_row(prepared_df)
            
            # Check if table is partitioned by checking for dt/DT column
            partition_field = self._detect_partition_field(prepared_df)
            has_partition_column = partition_field is not None
            
            if has_partition_column:
                # Partitioned table write
                if partition_value is None:
                    # Extract partition value from the data
                    # Assume all rows have the same dt value (common pattern)
                    partition_value = str(prepared_df[partition_field].iloc[0])
                    logger.debug(f"Extracted partition value from data: {partition_value}")
                
                logger.info(f"Writing to partitioned table with partition {partition_field}={partition_value}")
                
                # For partitioned tables, we need to write to the specific partition
                # Use insertInto with partition specification
                if mode == 'overwrite':
                    # Overwrite specific partition
                    spark_df.write \
                        .mode('overwrite') \
                        .insertInto(table_name, overwrite=True)
                    logger.info(f"Overwrote partition {partition_field}={partition_value} in table {table_name}")
                else:
                    # Append to partition
                    spark_df.write \
                        .mode('append') \
                        .insertInto(table_name)
                    logger.info(f"Appended to partition {partition_field}={partition_value} in table {table_name}")
                
            else:
                # Non-partitioned table write
                logger.info(f"Writing to non-partitioned table")
                
                if mode == 'overwrite':
                    spark_df.write \
                        .mode('overwrite') \
                        .insertInto(table_name, overwrite=True)
                    logger.info(f"Overwrote table {table_name}")
                else:
                    spark_df.write \
                        .mode('append') \
                        .insertInto(table_name)
                    logger.info(f"Appended to table {table_name}")
            
            logger.info(f"Successfully wrote {row_count} rows to {table_name}")
            return row_count
            
        except Exception as e:
            error_str = str(e)
            
            # Ignore java.io.FileNotFoundException for non-partitioned tables
            # This is a known Spark issue that doesn't affect data writing
            if 'java.io.FileNotFoundException' in error_str and not has_partition_column:
                logger.warning(f"Ignoring FileNotFoundException for non-partitioned table (data was written successfully)")
                logger.debug(f"FileNotFoundException details: {error_str}")
                logger.info(f"Successfully wrote {row_count} rows to {table_name} (FileNotFoundException ignored)")
                return row_count
            
            # For other errors, raise exception
            error_msg = f"Failed to write DataFrame to table {table_name}: {error_str}"
            logger.error(error_msg)
            raise HiveQueryError(error_msg) from e
    
    def close(self):
        """
        Close the Hive connection.
        
        Note:
            For Spark backend, this does NOT stop the SparkSession since it may be shared
            across multiple clients. For PyHive backend, this closes the connection.
        """
        if self._backend == 'spark' and self._spark:
            logger.info("Releasing reference to SparkSession (session remains active)")
            self._spark = None
        elif self._backend == 'pyhive' and self._pyhive_conn:
            logger.info("Closing PyHive connection")
            try:
                self._pyhive_conn.close()
            except Exception as e:
                logger.warning(f"Error closing PyHive connection: {e}")
            self._pyhive_conn = None
    
    def is_connected(self) -> bool:
        """
        Check if the Hive connection is active.
        
        Returns:
            bool: True if connection is active, False otherwise
        """
        if self._backend == 'spark':
            if self._spark is None:
                return False
            try:
                # Test connection by accessing sparkContext
                return self._spark.sparkContext._jsc is not None
            except Exception:
                return False
        elif self._backend == 'pyhive':
            if self._pyhive_conn is None:
                return False
            try:
                # Test connection by executing a simple query
                cursor = self._pyhive_conn.cursor()
                cursor.execute("SELECT 1")
                cursor.close()
                return True
            except Exception:
                return False
        else:
            return False
    
    def __enter__(self):
        """Context manager entry."""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.close()
        return False
