"""ODBC data source implementation."""

from __future__ import annotations

import os

from typing import TYPE_CHECKING, Any

import pyarrow as pa
import pyarrow.parquet as pq

from dependency_injector.wiring import Provide, inject

from data_exchange_agent.constants import config_keys, container_keys
from data_exchange_agent.constants.paths import build_actual_results_folder_path
from data_exchange_agent.data_sources.base import BaseDataSource
from data_exchange_agent.data_sources.odbc_output_converters import (
    register_output_converters,
)
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

    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


"""
The default number of rows to fetch in each batch.
"""
DEFAULT_ROWS_PER_BATCH = 50000

"""
The index of the name in the description of a column.
"""
ODBC_DESCRIPTION_NAME_INDEX = 0


class ODBCDataSource(BaseDataSource):
    """
    An ODBC data source implementation.

    This class provides a way to export data from an ODBC data source to a Parquet file.

    Attributes:
        statement (str): The SQL statement to execute
        results_folder_path (str): The path to the results folder
        batch_size (int): The number of rows to fetch in each batch

    """

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

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

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

    @property
    def batch_size(self) -> int:
        """The number of rows to fetch in each batch."""
        return self._batch_size

    @inject
    def __init__(
        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 ODBCDataSource.

        Args:
            engine (str): The database engine type
            statement (str): The SQL statement to execute
            results_folder_path (str): The path to the results folder
            base_file_name (str): The base file name
            suggested_batch_size (int | None): Suggested batch size for extraction.
                When None, defaults to DEFAULT_ROWS_PER_BATCH.
            database (str | None): Database name (kept for API compatibility with BCP).
                Not used directly by ODBCDataSource; see ``use_database_command``.
            use_database_command (str | None): Pre-built SQL command to switch the active
                database before executing the statement (e.g., ``USE [SampleStoreDB]``).
                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
        self._results_folder_path: str = (
            build_actual_results_folder_path() if results_folder_path is None else results_folder_path
        )
        self._base_file_name: str = base_file_name
        self._batch_size: int = suggested_batch_size if suggested_batch_size is not None else DEFAULT_ROWS_PER_BATCH
        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
        self.__connection_string: str = source_connections_config[engine].connection_string

    def export_data(self) -> bool:
        """
        Export data to a Parquet file.

        Returns:
            bool: True if the data was exported successfully, False otherwise

        Raises:
            Exception: If the SQL statement is not a read-only operation

        """
        # Check if the SQL statement is a read-only operation
        sql_command_type = get_permitted_sql_command_type(self.statement)
        if sql_command_type not in (
            SQLCommandType.SELECT,
            SQLCommandType.WITH,
            SQLCommandType.DESCRIBE,
            SQLCommandType.DESC,
            SQLCommandType.SHOW,
            SQLCommandType.EXPLAIN,
        ):
            raise Exception("The SQL statement is not a read-only operation.")

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

        # Register output converters for SQL Server types that pyodbc cannot
        # marshal natively (e.g. SQL_VARIANT, XML, DATETIMEOFFSET).  This
        # prevents ProgrammingError during cursor.fetchmany().
        register_output_converters(conn)

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

        try:
            # Export data to Parquet file
            self._export_sql_results_to_parquet(
                conn,
                self.statement,
                parquet_folder_path=self.results_folder_path,
                batch_size=self.batch_size,
            )
        except Exception as e:
            self.logger.error(f"Error exporting data to a Parquet file. Error: {e}")
            raise e
        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}")
            conn.close()
            raise e

    def _export_sql_results_to_parquet(
        self,
        conn: pyodbc.Connection,
        sql_query: str,
        parquet_folder_path: str = "data_chunks",
        batch_size: int = DEFAULT_ROWS_PER_BATCH,
    ) -> None:
        """
        Export SQL results to a Parquet file.

        Args:
            conn (pyodbc.Connection): The database connection
            sql_query (str): The SQL query to execute
            parquet_folder_path (str): The path to the output directory
            batch_size (int): The number of rows to fetch in each batch

        """
        cursor: pyodbc.Cursor | None = None
        try:
            # Execute SQL query
            cursor = conn.cursor()
            cursor.arraysize = batch_size
            self.logger.info(f"Start: Execution of query: {sql_query}...")
            cursor.execute(sql_query)
            self.logger.info(f"End: Execution of query: {sql_query}.")

            # Create output directory
            os.makedirs(parquet_folder_path, exist_ok=True)

            parquet_file = os.path.join(parquet_folder_path, f"{self.base_file_name}_001.parquet")
            self._write_chunks_to_single_parquet(cursor, parquet_file, batch_size=batch_size)
        finally:
            if cursor:
                cursor.close()

    def _write_chunks_to_single_parquet(
        self,
        cursor: pyodbc.Cursor,
        output_file: str,
        batch_size: int = DEFAULT_ROWS_PER_BATCH,
    ):
        """
        Write query results to a single Parquet file in chunks.

        Args:
            cursor (pyodbc.Cursor): The cursor to use for writing the results
            output_file (str): The path to the output file
            batch_size (int): The number of rows to fetch in each batch

        """
        self.logger.info(f"Start: Writing chunks to single parquet file {output_file}...")
        writer = None

        batch_index = 1
        try:
            while True:
                # Fetch many rows from the cursor
                self.logger.info(f"Start: Fetching batch #{batch_index} (rows: {batch_size})")
                rows = cursor.fetchmany(batch_size)
                self.logger.info(f"End: Fetching batch #{batch_index} (rows: {batch_size})")
                if not rows:
                    break

                # Convert to PyArrow Record Batch
                record_batch = self._build_pyarrow_record_batch(cursor.description or [], rows)

                if writer is None:
                    # Create writer with schema from first batch
                    writer = pq.ParquetWriter(output_file, record_batch.schema)

                # Write chunk as a row group
                writer.write_batch(record_batch)

                batch_index = batch_index + 1

        finally:
            if writer:
                writer.close()
        self.logger.info(f"End: Writing chunks to single parquet file {output_file}.")

    def _build_pyarrow_record_batch(self, description: list, rows: list) -> pa.RecordBatch:
        """
        Build a PyArrow record batch from a list of columns.

        Args:
            description (list): The descriptions of the columns.
            rows (list): The rows of data.

        Returns:
            pa.RecordBatch: The PyArrow record batch

        """
        column_names: list[str] = [desc[ODBC_DESCRIPTION_NAME_INDEX] for desc in description]
        arrays = [self._build_pyarrow_array(rows, column_index) for column_index in range(len(column_names))]
        return pa.RecordBatch.from_arrays(arrays, names=column_names)

    def _build_pyarrow_array(self, rows: list, column_index: int) -> pa.Array:
        """
        Build a PyArrow array that represents a column.

        ODBC returns native Python types, so we can directly convert to PyArrow arrays
        without the need for custom type mappers (unlike JDBC which returns Java objects).

        Args:
            rows (list): The rows of data.
            column_index (int): The index of the column to build

        Returns:
            pa.Array: The PyArrow array that represents the column.

        """
        column_values = [row[column_index] for row in rows]
        return pa.array(column_values)
