"""
Redshift executor factory.

Wraps the ``ConnectorRedshift`` from ``snowflake-data-validation`` to provide
connection management, SQL literal formatting, and stored-procedure execution.

Column types are captured from ``cursor.description`` and mapped to Snowflake
types through the framework's YAML type-mapping templates.
"""

from __future__ import annotations

import logging
from typing import Any

from snowflake.snowflake_data_validation.connector.connector_base import ConnectorBase
from snowflake.snowflake_data_validation.redshift.model.redshift_credentials_connection import (
    RedshiftCredentialsConnection,
)
from snowflake.snowflake_data_validation.redshift.connector.connector_factory_redshift import (
    RedshiftConnectorFactory,
)
from snowflake.snowflake_data_validation.utils.constants import Platform

from test_runner.common.database_dialect import DatabaseDialect
from test_runner.common.database_executor import (
    DatabaseExecutor,
    DatabaseExecutorFactory,
    LiteralFormatter,
)
from test_runner.common.models import (
    ColumnTypeMap,
    TestCaseResult,
    ValidationSteps,
)
from test_runner.common.type_mapping import base_type_name, load_datatypes_mapping
from test_runner.replay.isolation import RedshiftTransactionIsolation
from test_runner.replay.source_introspector import RedshiftIntrospector

import redshift_connector


LOGGER = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Literal formatter
# ---------------------------------------------------------------------------

class RedshiftLiteralFormatter(LiteralFormatter):
    """PostgreSQL/Redshift literal formatting: TRUE/FALSE for bools, single-quoted strings."""

    def format_literal(self, value: Any) -> str:
        if value is None:
            return "NULL"
        if isinstance(value, bool):
            return "TRUE" if value else "FALSE"
        if isinstance(value, (int, float)):
            return str(value)
        if isinstance(value, str):
            escaped = value.replace("'", "''")
            return f"'{escaped}'"
        escaped = str(value).replace("'", "''")
        return f"'{escaped}'"


# ---------------------------------------------------------------------------
# Executor
# ---------------------------------------------------------------------------

class RedshiftExecutor(DatabaseExecutor):
    """Executes SQL on Redshift via the data-validation ``ConnectorRedshift``.

    Uses the raw ``redshift_connector`` connection to execute queries and
    captures column type names from ``cursor.description``.

    Every execution is wrapped in a transaction (``BEGIN`` / ``ROLLBACK``).

    When ``output_parameters`` is declared, scalar OUT/INOUT values are
    read from the CALL result set and stored in ``result.output_params``.

    When ``fetch_stmts`` are provided (pre-rendered by Jinja templates),
    each is executed in sequence after the CALL and its result set is
    captured.  The executor doesn't know whether a fetch statement is a
    ``FETCH ALL`` (cursor) or ``SELECT *`` (table) — all dialect-specific
    SQL lives in the templates.
    """

    def __init__(
        self,
        connector: ConnectorBase,
        datatypes_mappings: ColumnTypeMap | None = None,
    ) -> None:
        self._connector = connector
        self._datatypes_mappings = datatypes_mappings
        # Enable autocommit so a failed statement doesn't poison subsequent
        # executions with "current transaction is aborted" errors.
        try:
            self._connector.connection.autocommit = True
        except Exception:
            pass

    @property
    def connector(self) -> ConnectorBase:
        return self._connector

    def close(self) -> None:
        self._connector.close()

    def execute(
        self,
        sql: str,
        steps: ValidationSteps | None = None,
        fetch_stmts: list[str] | None = None,
    ) -> TestCaseResult:
        output_params = (
            steps.output_parameters if steps else None
        )
        has_outputs = bool(output_params or fetch_stmts)

        result = TestCaseResult()
        conn = self._connector.connection
        cursor = conn.cursor()
        try:
            cursor.execute(sql)

            if has_outputs:
                if output_params:
                    self._extract_output_params(cursor, result, output_params)

                for stmt in (fetch_stmts or []):
                    cursor.execute(stmt)
                    self._append_result_set(cursor, result)
            else:
                self._append_result_set(cursor, result)
        except Exception as exc:
            self._try_rollback(cursor)
            result.success = False
            result.error = str(exc)
        finally:
            cursor.close()

        return result

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _try_rollback(cursor: redshift_connector.Cursor) -> None:
        """Issue a ROLLBACK to clear Redshift's aborted-transaction state."""
        try:
            cursor.execute("ROLLBACK")
        except Exception:
            pass

    def _append_result_set(self, cursor: redshift_connector.Cursor, result: TestCaseResult) -> None:
        """Read the current cursor result set and append it to *result*."""
        description = cursor.description
        if description is None:
            result.result_sets.append([])
            result.row_counts.append(0)
            result.column_types.append({})
            return

        column_names = [desc[0] for desc in description]
        rows = cursor.fetchall()
        row_dicts = [dict(zip(column_names, row)) for row in rows]
        result.result_sets.append(row_dicts)
        result.row_counts.append(len(row_dicts))
        result.column_types.append(self._describe_types(cursor))

    def _extract_output_params(
        self,
        cursor: redshift_connector.Cursor,
        result: TestCaseResult,
        param_names: list[str],
    ) -> None:
        """Read scalar OUT/INOUT values from the CALL result set.

        In Redshift, both OUT and INOUT parameters appear as columns in the
        single-row result set returned by the CALL.  This method reads that
        row, matches columns to *param_names* (case-insensitive), and
        populates ``result.output_params`` / ``result.output_param_types``.
        """
        description = cursor.description
        if description is None:
            return

        column_names = [desc[0] for desc in description]
        row = cursor.fetchone()
        if row is None:
            return

        col_types = self._describe_types(cursor)
        row_dict = dict(zip(column_names, row))

        lower_to_declared = {n.lower(): n for n in param_names}
        output_params: dict[str, Any] = {}
        output_param_types: dict[str, str] = {}

        for col_name, col_value in row_dict.items():
            declared = lower_to_declared.get(col_name.lower())
            if declared is not None:
                output_params[declared] = col_value

        for col_name, col_type in col_types.items():
            declared = lower_to_declared.get(col_name.lower())
            if declared is not None:
                output_param_types[declared] = col_type

        result.output_params = output_params
        result.output_param_types = output_param_types

    def _describe_types(
        self, cursor: redshift_connector.Cursor,
    ) -> ColumnTypeMap:
        """Map column OIDs to Snowflake type names via ``pg_catalog.pg_type``.

        Uses a separate cursor to query the catalog, matching the pattern
        used by the SQL Server executor (which queries
        ``sp_describe_first_result_set``).
        """
        description = cursor.description
        if not description:
            return {}

        oid_to_col: dict[int, list[str]] = {}
        for desc in description:
            col_name = desc[0].upper()
            type_code = desc[1]
            oid_to_col.setdefault(type_code, []).append(col_name)

        unique_oids = list(oid_to_col.keys())
        oid_list = ", ".join(str(o) for o in unique_oids)

        oid_name_map: dict[int, str] = {}
        try:
            meta_cursor = self._connector.connection.cursor()
            try:
                meta_cursor.execute(
                    f"SELECT oid, typname FROM pg_catalog.pg_type "
                    f"WHERE oid IN ({oid_list})"
                )
                for row in meta_cursor.fetchall():
                    oid_name_map[int(row[0])] = str(row[1])
            finally:
                meta_cursor.close()
        except Exception as exc:
            LOGGER.warning("pg_catalog.pg_type lookup failed: %s", exc)

        column_types: dict[str, str] = {}
        for desc in description:
            col_name = desc[0].upper()
            type_code = desc[1]
            type_name = oid_name_map.get(type_code, str(type_code))
            base_type = base_type_name(type_name)
            if self._datatypes_mappings:
                snowflake_type = self._datatypes_mappings.get(base_type, base_type)
            else:
                snowflake_type = base_type
            column_types[col_name] = snowflake_type
        return column_types


# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------

class RedshiftExecutorFactory(DatabaseExecutorFactory):
    """Factory for Redshift executors backed by data-validation connectors."""

    @property
    def dialect(self) -> DatabaseDialect:
        return DatabaseDialect.REDSHIFT

    def build_config(self, raw: dict[str, Any]) -> RedshiftCredentialsConnection:
        _AUTH_METHOD_MAP = {"standard": "credentials"}
        raw_mode = raw.get("mode") or raw.get("auth_method", "credentials")
        mode = _AUTH_METHOD_MAP.get(raw_mode, raw_mode)
        return RedshiftCredentialsConnection(
            mode=mode,
            host=raw.get("host") or raw.get("server_url", "localhost"),
            port=int(raw.get("port", 5439)),
            database=raw.get("database", ""),
            username=raw.get("username") or raw.get("user", ""),
            password=raw.get("password", ""),
        )

    def create_literal_formatter(self) -> RedshiftLiteralFormatter:
        return RedshiftLiteralFormatter()

    def create_executor(self, config: RedshiftCredentialsConnection) -> RedshiftExecutor:
        connector = RedshiftConnectorFactory.create_connector(config)
        datatypes_mappings = load_datatypes_mapping(Platform.REDSHIFT)
        return RedshiftExecutor(connector, datatypes_mappings)

    def create_isolation(self, **kwargs):
        return RedshiftTransactionIsolation()

    def create_introspector(self):
        return RedshiftIntrospector()
