"""
SQL Server executor factory.

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

Column types are captured via ``sp_describe_first_result_set`` which returns
the real SQL Server type names.  These are then 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, Sequence

from snowflake.snowflake_data_validation.connector.connector_base import ConnectorBase
from snowflake.snowflake_data_validation.sqlserver.model.sqlserver_credentials_connection import (
    SqlServerCredentialsConnection,
)
from snowflake.snowflake_data_validation.sqlserver.connector.connector_factory_sql_server import (
    SqlServerConnectorFactory,
)
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 SqlServerTransactionIsolation
from test_runner.replay.source_introspector import SqlServerIntrospector


LOGGER = logging.getLogger(__name__)

_PYODBC_TYPE_MAP: dict[type, str] = {
    str: "varchar",
    int: "int",
    float: "float",
    bool: "bit",
    bytes: "varbinary",
    decimal.Decimal: "decimal",
    datetime.datetime: "datetime",
    datetime.date: "date",
    datetime.time: "time",
}


def _parse_sp_describe_rows(
    meta_rows: Sequence[Any],
    datatypes_mappings: ColumnTypeMap | None = None,
) -> ColumnTypeMap:
    """Build column types from ``sp_describe_first_result_set`` output.

    Each row contains the column name at index 2 and the SQL Server type
    name (e.g. ``"int"``, ``"nvarchar(100)"``) at index 5.  The base type
    name is looked up in *datatypes_mappings* to produce the Snowflake type.

    The base type name is looked up in *datatypes_mappings* to produce
    the Snowflake type.
    """
    column_types: dict[str, str] = {}
    for row in meta_rows:
        col_name: str = row[2]
        system_type_name: str = row[5]
        base_type = base_type_name(system_type_name)
        if datatypes_mappings:
            snowflake_type = datatypes_mappings.get(base_type, base_type)
        else:
            snowflake_type = base_type
        column_types[col_name.upper()] = snowflake_type
    return column_types


# ---------------------------------------------------------------------------
# Literal formatter (SRP: separated from factory)
# ---------------------------------------------------------------------------

class SqlServerLiteralFormatter(LiteralFormatter):
    """T-SQL literal formatting: 1/0 for bools, single-quoted strings."""

    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 SqlServerExecutor(DatabaseExecutor):
    """Executes SQL on SQL Server via the data-validation ``ConnectorSqlServer``.

    Uses the raw pyodbc connection to execute queries and captures real
    column type names via ``sp_describe_first_result_set``.
    """

    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:
        result = TestCaseResult()
        try:
            outputs = steps.output_parameters if steps else None

            with self._connector.connection.cursor() as cursor:
                cursor.execute(sql)

                if outputs:
                    # When outputs exist the Jinja template appends a
                    # ``SELECT @var AS [name], …`` as the LAST statement in the
                    # batch.  Collect every result set; the last one is the
                    # output-param row and anything before it belongs to the
                    # procedure's own result sets.
                    all_sets: list[tuple[list[str], list, Any]] = []
                    if cursor.description:
                        cols = [d[0] for d in cursor.description]
                        desc = cursor.description
                        all_sets.append((cols, cursor.fetchall(), desc))
                    while cursor.nextset():
                        if cursor.description:
                            cols = [d[0] for d in cursor.description]
                            desc = cursor.description
                            all_sets.append((cols, cursor.fetchall(), desc))

                    if all_sets:
                        out_cols, out_rows, out_desc = all_sets[-1]
                        if out_rows:
                            result.output_params = dict(
                                zip(out_cols, out_rows[0])
                            )
                            result.output_param_types = (
                                self._types_from_description(out_desc)
                            )
                        for col_names, rows, _ in all_sets[:-1]:
                            row_dicts = [dict(zip(col_names, r)) for r in rows]
                            result.result_sets.append(row_dicts)
                            result.row_counts.append(len(row_dicts))
                            result.column_types.append({})
                else:
                    if cursor.description:
                        column_names = [d[0] for d in cursor.description]
                        rows = cursor.fetchall()
                        row_dicts = [dict(zip(column_names, r)) for r in rows]
                        result.result_sets = [row_dicts]
                        result.row_counts = [len(row_dicts)]
                        result.column_types = [{}]

                # Track how many result sets came from the CALL itself so
                # the post-loop ``sp_describe_first_result_set`` only covers
                # those (fetch_stmts results get their types inline).
                pre_fetch_count = len(result.result_sets)

                # Execute fetch statements (output.tables / output.cursors).
                # Each statement (e.g. ``SELECT * FROM <temp_table>``) is run
                # after the main CALL and its result set is appended to the
                # capture result — matching the Redshift executor behaviour.
                for stmt in (fetch_stmts or []):
                    cursor.execute(stmt)
                    self._append_fetch_result(cursor, result)

            # Resolve column types for the CALL's own result sets via
            # sp_describe_first_result_set.  Fetch-stmt results already
            # have their types populated by _append_fetch_result.
            for i in range(pre_fetch_count):
                result.column_types[i] = self._describe_types(sql)

        except Exception as exc:
            result.success = False
            result.error = str(exc)
        return result

    def _append_fetch_result(self, cursor: Any, result: TestCaseResult) -> None:
        """Read the current cursor result set and append it to *result*."""
        if cursor.description is None:
            result.result_sets.append([])
            result.row_counts.append(0)
            result.column_types.append({})
            return
        column_names = [d[0] for d in cursor.description]
        rows = cursor.fetchall()
        row_dicts = [dict(zip(column_names, r)) for r in rows]
        result.result_sets.append(row_dicts)
        result.row_counts.append(len(row_dicts))
        result.column_types.append(self._types_from_description(cursor.description))

    def _types_from_description(
        self, description: Sequence[Any],
    ) -> ColumnTypeMap:
        """Build column types from a pyodbc ``cursor.description``.

        Maps Python DB-API type objects to SQL Server type names, then
        through ``_datatypes_mappings`` to produce Snowflake types.
        """
        import pyodbc

        col_types: dict[str, str] = {}
        for desc in description:
            col_name = desc[0].upper()
            type_code = desc[1]
            sql_type = _PYODBC_TYPE_MAP.get(type_code, "varchar")
            if self._datatypes_mappings:
                snowflake_type = self._datatypes_mappings.get(
                    sql_type, sql_type.upper(),
                )
            else:
                snowflake_type = sql_type.upper()
            col_types[col_name] = snowflake_type
        return col_types

    def _describe_types(self, sql: str) -> ColumnTypeMap:
        """Get real column type names via ``sp_describe_first_result_set``."""
        try:
            with self._connector.connection.cursor() as cursor:
                cursor.execute(
                    "EXEC sp_describe_first_result_set @tsql = ?", sql,
                )
                meta_rows = cursor.fetchall()
            return _parse_sp_describe_rows(meta_rows, self._datatypes_mappings)
        except Exception as exc:
            LOGGER.warning("sp_describe_first_result_set failed: %s", exc)
            return {}


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

class SqlServerExecutorFactory(DatabaseExecutorFactory):
    """Factory for SQL Server executors backed by data-validation connectors."""

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

    @staticmethod
    def _to_yes_no(value: Any, default: str) -> str:
        """Normalise bool/string values to the 'yes'/'no' expected by ODBC Driver 18."""
        if isinstance(value, bool):
            return "yes" if value else "no"
        if isinstance(value, str):
            return "yes" if value.lower() in ("yes", "true", "1") else "no"
        return default

    def build_config(self, raw: dict[str, Any]) -> SqlServerCredentialsConnection:
        # Accept both SCAI-generated keys (server_url, user, auth_method) and
        # the canonical test-runner keys (host, username, mode).
        # SCAI uses auth_method="standard" which maps to sql_auth.
        _AUTH_METHOD_MAP = {"standard": "sql_auth", "credentials": "sql_auth"}
        raw_mode = raw.get("mode") or raw.get("auth_method", "sql_auth")
        mode = _AUTH_METHOD_MAP.get(raw_mode, raw_mode)
        return SqlServerCredentialsConnection(
            mode=mode,
            host=raw.get("host") or raw.get("server_url", "localhost"),
            port=int(raw.get("port", 1433)),
            database=raw.get("database", ""),
            username=raw.get("username") or raw.get("user"),
            password=raw.get("password"),
            client_id=raw.get("client_id"),
            client_secret=raw.get("client_secret"),
            tenant_id=raw.get("tenant_id"),
            trust_server_certificate=self._to_yes_no(
                raw.get("trust_server_certificate"), "no"
            ),
            encrypt=self._to_yes_no(raw.get("encrypt"), "yes"),
        )

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

    def create_executor(self, config: SqlServerCredentialsConnection) -> SqlServerExecutor:
        connector = SqlServerConnectorFactory.create_connector(config)
        datatypes_mappings = load_datatypes_mapping(Platform.SQLSERVER)
        return SqlServerExecutor(connector, datatypes_mappings)

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

    def create_introspector(self):
        return SqlServerIntrospector()
