"""
Redshift UNLOAD data source implementation.

This module provides a specialized data source for executing Redshift UNLOAD
commands. Unlike traditional data sources, UNLOAD operations write directly
to S3 and do not return result sets to fetch.
"""

from __future__ import annotations

import re

from typing import TYPE_CHECKING, Any

from dependency_injector.wiring import Provide, inject

from data_exchange_agent.constants import config_keys, container_keys
from data_exchange_agent.data_sources.base import BaseDataSource
from data_exchange_agent.data_sources.sql_command_type import SQLCommandType
from data_exchange_agent.data_sources.sql_parser import get_permitted_sql_command_type
from data_exchange_agent.utils.sf_logger import SFLogger


if TYPE_CHECKING:
    import pyodbc  # pylint: disable=import-error

    from data_exchange_agent.config.manager import ConfigManager
    from data_exchange_agent.config.sections.connections.odbc.base import (
        BaseODBCConnectionConfig,
    )


def _get_pyodbc() -> Any:
    """Lazily import pyodbc to avoid loading unixODBC at module import time."""
    import pyodbc

    return pyodbc


# Pattern to match IAM_ROLE 'arn:...' (single or double quoted)
_IAM_ROLE_PATTERN = re.compile(
    r"(IAM_ROLE\s+)['\"]([^'\"]+)['\"]",
    re.IGNORECASE,
)


def _sanitize_unload_statement(statement: str) -> str:
    """
    Redact sensitive credentials from an UNLOAD statement for safe logging.

    Replaces the IAM role ARN with a masked placeholder so that role ARNs
    are never written to log files.

    Args:
        statement: The raw UNLOAD SQL statement.

    Returns:
        The statement with IAM role ARN redacted.

    """
    return _IAM_ROLE_PATTERN.sub(r"\1'***'", statement)


class UnloadDataSource(BaseDataSource):
    """
    Data source for Redshift UNLOAD operations.

    This class executes Redshift UNLOAD commands which write data directly
    to S3. Unlike standard data sources, UNLOAD does not return result sets,
    so this implementation:
    - Executes the UNLOAD command
    - Does NOT fetch any data locally
    - Does NOT create Parquet files
    - Returns success after command execution

    The UNLOAD command itself contains the S3 destination, IAM role, and
    file format configuration. The data is written directly from Redshift
    to S3, bypassing the agent entirely.

    Attributes:
        statement (str): The UNLOAD SQL statement to execute
        results_folder_path (str): Unused for UNLOAD (kept for interface compatibility)
        base_file_name (str): Unused for UNLOAD (kept for interface compatibility)

    """

    @property
    def statement(self) -> str:
        """The UNLOAD statement to execute."""
        return self._statement

    @property
    def results_folder_path(self) -> str:
        """The path to the results folder (unused for UNLOAD)."""
        return self._results_folder_path

    @property
    def base_file_name(self) -> str:
        """The base file name (unused for UNLOAD)."""
        return self._base_file_name

    @inject
    def __init__(  # pylint: disable=unused-argument
        self,
        engine: str,
        statement: str,
        results_folder_path: str | None = None,
        base_file_name: str = "result",
        suggested_batch_size: int | None = None,
        database: str | None = None,
        use_database_command: str | None = None,
        logger: SFLogger = Provide[container_keys.SF_LOGGER],
        program_config: ConfigManager = Provide[container_keys.PROGRAM_CONFIG],
    ) -> None:
        """
        Initialize a new UnloadDataSource.

        Args:
            engine (str): The database engine type (should be 'redshift')
            statement (str): The UNLOAD SQL statement to execute
            results_folder_path (str): Unused for UNLOAD operations
            base_file_name (str): Unused for UNLOAD operations
            suggested_batch_size (int | None): Unused for UNLOAD operations
            database (str | None): Database name (kept for API compatibility)
            use_database_command (str | None): Pre-built SQL command to switch the active
                database before executing the statement (e.g., ``USE database_name``).
                Generated by the orchestrator via the platform's query builder.
                When None, no database switching is performed.
            logger (SFLogger): The logger instance
            program_config (ConfigManager): The program configuration manager

        """
        self.logger: SFLogger = logger
        self._statement: str = statement
        # UNLOAD writes directly to S3; no local results folder is needed.
        self._results_folder_path: str = results_folder_path or ""
        self._base_file_name: str = base_file_name
        self._use_database_command: str | None = use_database_command
        source_connections_config: dict[str, BaseODBCConnectionConfig] = program_config[
            config_keys.CONNECTIONS__SOURCE
        ]

        # Initialize ODBC connection settings for Redshift
        self.__connection_string: str = source_connections_config[
            engine
        ].connection_string

    def export_data(self) -> bool:
        """
        Execute the UNLOAD command.

        This method connects to Redshift and executes the UNLOAD statement.
        UNLOAD writes data directly to S3 and does not return any result set,
        so we simply execute the command and return success.

        Returns:
            bool: True if the UNLOAD command executed successfully

        Raises:
            Exception: If the statement is not an UNLOAD command or if execution fails

        """
        # Validate that this is an UNLOAD command
        sql_command_type = get_permitted_sql_command_type(self.statement)
        if sql_command_type != SQLCommandType.UNLOAD:
            raise Exception(
                f"UnloadDataSource can only execute UNLOAD commands, got: {sql_command_type}"
            )

        pyodbc = _get_pyodbc()
        conn: pyodbc.Connection | None = None
        try:
            # Create a connection to Redshift
            self.logger.info("Connecting to Redshift for UNLOAD operation")
            conn = pyodbc.connect(self.__connection_string)
        except Exception as e:
            self.logger.error(f"Error creating a connection to Redshift. Error: {e}")
            raise

        # Switch to the specified database if a USE command was provided by the orchestrator
        if self._use_database_command:
            self._switch_database(conn)

        try:
            # Execute the UNLOAD command
            # UNLOAD writes directly to S3 and does not return rows
            self.logger.info(
                f"Executing UNLOAD command: {_sanitize_unload_statement(self.statement)}"
            )
            cursor = conn.cursor()
            cursor.execute(self.statement)
            cursor.close()
            self.logger.info(
                "UNLOAD command executed successfully. Data written to S3."
            )
        except Exception as e:
            self.logger.error(f"Error executing UNLOAD command. Error: {e}")
            raise
        finally:
            if conn:
                conn.close()

        return True

    def _switch_database(self, conn: pyodbc.Connection) -> None:
        """
        Execute the pre-built USE database command on the connection.

        The command is generated by the orchestrator's platform query builder
        and passed through the task payload, so this method is completely
        dialect-agnostic.

        Args:
            conn: Active ODBC connection

        Raises:
            Exception: If the USE command fails.

        """
        assert self._use_database_command is not None  # Caller guarantees this

        try:
            self.logger.info(
                f"Executing database switch command: {self._use_database_command}"
            )
            cursor = conn.cursor()
            cursor.execute(self._use_database_command)
            cursor.close()
            self.logger.info(f"Successfully executed: {self._use_database_command}")
        except Exception as e:
            self.logger.error(
                f"Error executing database switch command '{self._use_database_command}'. Error: {e}"
            )
            raise
