"""
Connection testing utilities.

This module provides functionality to test database connections by executing
a simple SELECT 1 query to verify connectivity.
"""

from __future__ import annotations

import traceback

from dataclasses import dataclass
from enum import Enum
from typing import TYPE_CHECKING, Any


if TYPE_CHECKING:
    from data_exchange_agent.config.manager import ConfigManager


class ConnectionTestResult(Enum):
    """Result status for connection tests."""

    SUCCESS = "success"
    FAILURE = "failure"
    SKIPPED = "skipped"


@dataclass
class ConnectionTestOutcome:
    """Result of a connection test."""

    name: str
    status: ConnectionTestResult
    message: str
    error_details: str | None = None

    @property
    def is_success(self) -> bool:
        """Check if the test was successful."""
        return self.status == ConnectionTestResult.SUCCESS


# Alias for backwards compatibility
TestResult = ConnectionTestOutcome


class ConnectionTester:
    """
    Test database connections by executing SELECT 1.

    This class provides methods to test ODBC source connections and
    Snowflake target connections.
    """

    def __init__(self, config: ConfigManager) -> None:
        """
        Initialize the connection tester.

        Args:
            config: The configuration manager with connection settings.

        """
        self._config = config

    def test_all_connections(self) -> tuple[list[TestResult], bool]:
        """
        Test all configured connections.

        Returns:
            Tuple of (list of test results, all_passed boolean)

        """
        results: list[TestResult] = []
        all_passed = True

        # Test source connections (ODBC)
        source_results = self.test_source_connections()
        results.extend(source_results)
        if any(not r.is_success for r in source_results):
            all_passed = False

        # Test target connections (Snowflake)
        target_results = self.test_target_connections()
        results.extend(target_results)
        if any(r.status == ConnectionTestResult.FAILURE for r in target_results):
            all_passed = False

        return results, all_passed

    def test_source_connections(self) -> list[TestResult]:
        """
        Test all source (ODBC) connections.

        Returns:
            List of test results for each source connection.

        """
        from data_exchange_agent.constants import config_keys

        results: list[TestResult] = []
        source_connections = self._config.get(config_keys.CONNECTIONS__SOURCE, {})

        if not source_connections:
            return results

        for conn_name, conn_config in source_connections.items():
            result = self._test_odbc_connection(conn_name, conn_config)
            results.append(result)

        return results

    def test_target_connections(self) -> list[TestResult]:
        """
        Test all target (Snowflake) connections.

        Returns:
            List of test results for each target connection.

        """
        from data_exchange_agent.constants import config_keys

        results: list[TestResult] = []
        target_connections = self._config.get(config_keys.CONNECTIONS__TARGET, {})

        if not target_connections:
            return results

        for conn_name, conn_config in target_connections.items():
            result = self._test_snowflake_connection(conn_name, conn_config)
            if result.status != ConnectionTestResult.SKIPPED:
                results.append(result)

        return results

    def _test_odbc_connection(self, conn_name: str, conn_config: Any) -> TestResult:
        """
        Test an ODBC connection by executing SELECT 1.

        Args:
            conn_name: Connection name for display
            conn_config: The connection configuration

        Returns:
            TestResult with success/failure status

        """
        try:
            import pyodbc

            conn = pyodbc.connect(conn_config.connection_string)
            try:
                cursor = conn.cursor()
                cursor.execute("SELECT 1")
                cursor.fetchone()
                cursor.close()
                return TestResult(
                    name=conn_name,
                    status=ConnectionTestResult.SUCCESS,
                    message="Connection OK",
                )
            finally:
                conn.close()

        except Exception as e:
            full_trace = traceback.format_exc()
            return TestResult(
                name=conn_name,
                status=ConnectionTestResult.FAILURE,
                message=str(e),
                error_details=full_trace,
            )

    def _test_snowflake_connection(
        self, conn_name: str, conn_config: Any
    ) -> TestResult:
        """
        Test a Snowflake connection by executing SELECT 1.

        Args:
            conn_name: Connection name for display
            conn_config: The connection configuration

        Returns:
            TestResult with success/failure/skipped status

        """
        from data_exchange_agent.config.sections.connections.cloud_storages.snowflake import (
            SnowflakeConnectionConfig,
            SnowflakeConnectionNameConfig,
            SnowflakeConnectionPasswordConfig,
            SnowflakeExtendedBaseConnectionConfig,
        )

        # Check if this is a Snowflake connection
        if not isinstance(conn_config, SnowflakeConnectionConfig):
            return TestResult(
                name=conn_name,
                status=ConnectionTestResult.SKIPPED,
                message="Not a Snowflake connection",
            )

        try:
            import snowflake.connector

            # Build connection parameters based on config type
            if isinstance(conn_config, SnowflakeConnectionNameConfig):
                conn = snowflake.connector.connect(
                    connection_name=conn_config.connection_name
                )
            elif isinstance(conn_config, SnowflakeConnectionPasswordConfig):
                conn = snowflake.connector.connect(
                    account=conn_config.account,
                    user=conn_config.user,
                    password=conn_config.password,
                    warehouse=conn_config.warehouse,
                    database=conn_config.database,
                    schema=conn_config.schema,
                    role=getattr(conn_config, "role", None),
                )
            elif isinstance(conn_config, SnowflakeExtendedBaseConnectionConfig):
                conn = snowflake.connector.connect(
                    account=conn_config.account,
                    user=conn_config.user,
                    authenticator="externalbrowser",
                    warehouse=conn_config.warehouse,
                    database=conn_config.database,
                    schema=conn_config.schema,
                    role=getattr(conn_config, "role", None),
                )
            else:
                return TestResult(
                    name=conn_name,
                    status=ConnectionTestResult.SKIPPED,
                    message="Unsupported Snowflake config type",
                )

            try:
                cursor = conn.cursor()
                cursor.execute("SELECT 1")
                cursor.fetchone()
                cursor.close()
                return TestResult(
                    name=conn_name,
                    status=ConnectionTestResult.SUCCESS,
                    message="Connection OK",
                )
            finally:
                conn.close()

        except ImportError:
            return TestResult(
                name=conn_name,
                status=ConnectionTestResult.FAILURE,
                message="snowflake-connector-python not installed",
            )
        except Exception as e:
            full_trace = traceback.format_exc()
            return TestResult(
                name=conn_name,
                status=ConnectionTestResult.FAILURE,
                message=str(e),
                error_details=full_trace,
            )


def print_test_results(results: list[TestResult], all_passed: bool) -> None:
    """
    Print test results to console with colored output.

    Args:
        results: List of test results to print
        all_passed: Whether all tests passed

    """
    green = "\033[92m"
    red = "\033[91m"
    reset = "\033[0m"

    for result in results:
        if result.status == ConnectionTestResult.SUCCESS:
            print(f"  {green}[PASS]{reset} {result.name}: {result.message}")
        elif result.status == ConnectionTestResult.FAILURE:
            print(f"  {red}[FAIL]{reset} {result.name}: {result.message}")
            if result.error_details:
                for line in result.error_details.strip().split("\n"):
                    print(f"      {line}")
