"""
Teradata executor factory.

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

Column types are derived from ``cursor.description`` (which uses Python type
classes as type codes in ``teradatasql``) and mapped to Snowflake types through
the framework's YAML type-mapping templates.
"""

from __future__ import annotations

import datetime
import decimal
import logging
from typing import Any

from snowflake.snowflake_data_validation.connector.connector_base import ConnectorBase
from snowflake.snowflake_data_validation.teradata.model.teradata_credentials_connection import (
    TeradataCredentialsConnection,
)
from snowflake.snowflake_data_validation.teradata.connector.connector_factory_teradata import (
    TeradataConnectorFactory,
)
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 TeradataTransactionIsolation
from test_runner.replay.source_introspector import TeradataIntrospector


LOGGER = logging.getLogger(__name__)

# teradatasql uses Python type classes as cursor.description type codes.
_PYTHON_TYPE_TO_TD: dict[type, str] = {
    int: "integer",
    float: "float",
    str: "varchar",
    bytes: "varbyte",
    decimal.Decimal: "decimal",
    datetime.datetime: "timestamp",
    datetime.date: "date",
    datetime.time: "time",
    datetime.timedelta: "interval",
}


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

class TeradataLiteralFormatter(LiteralFormatter):
    """Teradata SQL literal formatting.

    Booleans are rendered as 0/1 (Teradata has no native boolean type;
    BYTEINT is conventionally used).
    """

    def format_literal(self, value: Any) -> str:
        if value is None:
            return "NULL"
        if isinstance(value, bool):
            return "1" if value else "0"
        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 TeradataExecutor(DatabaseExecutor):
    """Executes SQL on Teradata via the data-validation ``ConnectorTeradata``.

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

    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.
    """

    def __init__(
        self,
        connector: ConnectorBase,
        datatypes_mappings: ColumnTypeMap | None = None,
    ) -> None:
        self._connector = connector
        self._datatypes_mappings = datatypes_mappings

    @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._collect_all_result_sets(cursor, result)
        except Exception as exc:
            result.success = False
            result.error = str(exc)
        finally:
            cursor.close()

        return result

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

    def _collect_all_result_sets(
        self, cursor: Any, result: TestCaseResult,
    ) -> None:
        """Collect every non-empty result set from *cursor*.

        Teradata procedures declared with ``DYNAMIC RESULT SETS`` return
        cursor-based result sets that appear as subsequent result sets on
        the DB-API cursor after the initial ``execute``.  The first result
        set from the CALL itself is typically empty — some driver versions
        report ``description is None`` while others provide a description
        with zero rows.  We skip both cases and only keep result sets that
        actually contain rows.
        """
        try:
            has_more = True
            while has_more:
                description = cursor.description
                if description is not None:
                    column_names = [desc[0] for desc in description]
                    col_types = self._describe_types(cursor)
                    rows = cursor.fetchall()
                    if rows:
                        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(col_types)
                has_more = cursor.nextset()
        except Exception:
            LOGGER.debug("nextset() iteration ended", exc_info=True)

    def _append_result_set(self, cursor: Any, 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: Any,
        result: TestCaseResult,
        param_names: list[str],
    ) -> None:
        """Read scalar OUT/INOUT values from the CALL result set.

        In Teradata, output parameters from stored procedures appear as
        columns in the result set returned by the CALL.
        """
        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: Any) -> ColumnTypeMap:
        """Map cursor.description type codes to Snowflake type names.

        ``teradatasql`` uses Python type classes (``int``, ``str``,
        ``datetime.datetime``, etc.) as type codes.  These are mapped
        to Teradata type names and then through ``_datatypes_mappings``
        to produce the Snowflake type.
        """
        description = cursor.description
        if not description:
            return {}

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


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

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

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

    def build_config(self, raw: dict[str, Any]) -> TeradataCredentialsConnection:
        return TeradataCredentialsConnection(
            host=raw.get("host") or raw.get("server_url", "localhost"),
            database=raw.get("database", ""),
            username=raw.get("username") or raw.get("user", ""),
            password=raw.get("password", ""),
        )

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

    def create_executor(self, config: TeradataCredentialsConnection) -> TeradataExecutor:
        connector = TeradataConnectorFactory.create_connector(config)
        datatypes_mappings = load_datatypes_mapping(Platform.TERADATA)
        return TeradataExecutor(connector, datatypes_mappings)

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

    def create_introspector(self):
        return TeradataIntrospector()
