"""
Main data validation handler for the DEA.

Routes incoming ``data_validation`` tasks to the appropriate validation
strategy (schema, metrics, row-by-row, cell-by-cell) and provides shared
utilities for query execution and result I/O.
"""

import csv
import os
import shutil
import tempfile

from concurrent.futures import ThreadPoolExecutor
from typing import TYPE_CHECKING, Any

import pandas as pd

from data_exchange_agent.constants import task_keys
from data_exchange_agent.data_sources.odbc_output_converters import (
    register_output_converters,
)
from data_exchange_agent.data_sources.sf_connection import SnowflakeDataSource
from data_exchange_agent.tasks.data_validation.helpers import (
    CSV_DELIMITER,
    PAYLOAD_COLUMN_MAPPINGS,
    PAYLOAD_DATATYPES_MAPPINGS,
    PAYLOAD_ROW_VALIDATION_MODE,
    PAYLOAD_SOURCE_PLATFORM,
    PAYLOAD_SOURCE_QUERY,
    PAYLOAD_TABLE_METADATA_ID,
    PAYLOAD_TABLE_NAME,
    PAYLOAD_TARGET_ID,
    PAYLOAD_TARGET_QUERY,
    PAYLOAD_TOLERANCE,
    PAYLOAD_USE_DATABASE_COMMAND,
    PAYLOAD_VALIDATION_TYPE,
    PAYLOAD_WORKFLOW_ID,
    VALIDATION_TYPE_METRICS,
    VALIDATION_TYPE_ROW,
    VALIDATION_TYPE_SCHEMA,
    TaskResult,
    adapt_sdv_result,
    build_table_context,
    load_datatypes_mappings,
    resolve_platform,
)
from data_exchange_agent.utils.sf_logger import SFLogger
from snowflake.snowflake_data_validation.public import (
    QueryExecutionResult,
    validate_metrics,
    validate_schema,
)


if TYPE_CHECKING:
    import pyodbc

    from data_exchange_agent.config.sections.connections.odbc.base import (
        BaseODBCConnectionConfig,
    )


class DataValidationHandler:
    """Handle ``data_validation`` tasks."""

    def __init__(
        self,
        logger: SFLogger,
        snowflake_datasource: SnowflakeDataSource,
        source_connections_config: dict[str, "BaseODBCConnectionConfig"],
    ) -> None:
        """Initialize the handler with shared dependencies."""
        self.logger = logger
        self.snowflake_datasource = snowflake_datasource
        self.source_connections_config = source_connections_config

        from data_exchange_agent.tasks.data_validation.cell_validation import (
            CellValidationHandler,
        )
        from data_exchange_agent.tasks.data_validation.row_validation import (
            RowValidationHandler,
        )

        self._row_validator = RowValidationHandler(self)
        self._cell_validator = CellValidationHandler(self)

    # ------------------------------------------------------------------
    # Task routing
    # ------------------------------------------------------------------

    def handle(self, task: dict[str, Any]) -> TaskResult:
        """Process a data_validation task."""
        payload = task.get(task_keys.PAYLOAD, task)
        validation_type = payload.get(PAYLOAD_VALIDATION_TYPE, "")
        table_name = payload.get(PAYLOAD_TABLE_NAME, "unknown")
        workflow_id = payload.get(PAYLOAD_WORKFLOW_ID, 0)
        table_metadata_id = payload.get(PAYLOAD_TABLE_METADATA_ID, 0)
        source_query = payload.get(PAYLOAD_SOURCE_QUERY, "")
        target_query = payload.get(PAYLOAD_TARGET_QUERY, "")
        target_id = payload.get(PAYLOAD_TARGET_ID, "")
        column_mappings = payload.get(PAYLOAD_COLUMN_MAPPINGS) or {}
        datatypes_mappings = payload.get(PAYLOAD_DATATYPES_MAPPINGS) or {}
        tolerance = payload.get(PAYLOAD_TOLERANCE, 0.001)
        source_platform_str = payload.get(PAYLOAD_SOURCE_PLATFORM, "")

        self.logger.info(
            f"Processing data_validation task: type={validation_type}, "
            f"table={table_name}, workflow={workflow_id}"
        )

        try:
            if validation_type == VALIDATION_TYPE_ROW:
                row_mode = payload.get(PAYLOAD_ROW_VALIDATION_MODE, "row")
                if row_mode == "cell":
                    return self._cell_validator.handle(task, payload)
                return self._row_validator.handle(task, payload)

            if validation_type not in (VALIDATION_TYPE_SCHEMA, VALIDATION_TYPE_METRICS):
                raise ValueError(f"Unknown validation_type: {validation_type}")

            with ThreadPoolExecutor(max_workers=2) as pool:
                source_future = pool.submit(self._run_source_query, source_query, task)
                target_future = pool.submit(self._run_target_query, target_query)
                source_df = source_future.result()
                target_df = target_future.result()

            if "COLUMN_VALIDATED" in source_df.columns:
                source_df = source_df.copy()
                source_df["COLUMN_VALIDATED"] = source_df[
                    "COLUMN_VALIDATED"
                ].str.upper()
            if "COLUMN_VALIDATED" in target_df.columns:
                target_df = target_df.copy()
                target_df["COLUMN_VALIDATED"] = target_df[
                    "COLUMN_VALIDATED"
                ].str.upper()

            source_platform = resolve_platform(source_platform_str)
            table_ctx = build_table_context(source_platform, table_name)

            if not datatypes_mappings:
                datatypes_mappings = load_datatypes_mappings(source_platform)

            if validation_type == VALIDATION_TYPE_SCHEMA:
                result_df = self._validate_schema_sdv(
                    source_df,
                    target_df,
                    table_ctx,
                    column_mappings,
                    datatypes_mappings,
                    table_name,
                )
            else:
                result_df = self._validate_metrics_sdv(
                    source_df,
                    target_df,
                    table_ctx,
                    column_mappings,
                    tolerance,
                    table_name,
                )

            csv_path = self._write_result_csv(result_df, workflow_id, table_metadata_id)
            self._upload_to_stage(csv_path, target_id)
            self._cleanup_temp_dir(os.path.dirname(csv_path))

            self.logger.info(
                f"Validation complete for {table_name}: "
                f"{len(result_df)} result rows uploaded to {target_id}"
            )

            return {
                task_keys.STATUS: "success",
                task_keys.DETAILS: f"Validated {table_name} ({validation_type})",
                task_keys.TOTAL: len(result_df),
            }

        except Exception as e:
            self.logger.error(
                f"Data validation failed for {table_name}: {e}",
                exception=e,
            )
            return {
                task_keys.STATUS: "error",
                task_keys.DETAILS: str(e),
            }

    # ------------------------------------------------------------------
    # L1/L2 validation via SDV
    # ------------------------------------------------------------------

    def _validate_schema_sdv(
        self,
        source_df: pd.DataFrame,
        target_df: pd.DataFrame,
        table_ctx: Any,
        column_mappings: dict[str, str],
        datatypes_mappings: dict[str, str],
        table_name: str,
    ) -> pd.DataFrame:
        """Delegate schema validation to the SDV library."""
        if "COLUMN_VALIDATED" in source_df.columns:
            source_df = source_df.copy()
            source_df["COLUMN_VALIDATED"] = source_df["COLUMN_VALIDATED"].str.upper()
        if "COLUMN_VALIDATED" in target_df.columns:
            target_df = target_df.copy()
            target_df["COLUMN_VALIDATED"] = target_df["COLUMN_VALIDATED"].str.upper()

        src_result = QueryExecutionResult(dataframe=source_df)
        tgt_result = QueryExecutionResult(dataframe=target_df)

        validation = validate_schema(
            source_query_result=src_result,
            target_query_result=tgt_result,
            table_context=table_ctx,
            column_mappings=column_mappings,
            datatypes_mappings=datatypes_mappings,
        )

        self.logger.info(
            f"SDV schema validation for {table_name}: "
            f"is_valid={validation.is_valid}, rows={len(validation.results_dataframe)}"
        )
        return adapt_sdv_result(validation.results_dataframe, table_name)

    def _validate_metrics_sdv(
        self,
        source_df: pd.DataFrame,
        target_df: pd.DataFrame,
        table_ctx: Any,
        column_mappings: dict[str, str],
        tolerance: float,
        table_name: str,
    ) -> pd.DataFrame:
        """Delegate metrics validation to the SDV library."""
        src_result = QueryExecutionResult(dataframe=source_df)
        tgt_result = QueryExecutionResult(dataframe=target_df)

        validation = validate_metrics(
            source_query_result=src_result,
            target_query_result=tgt_result,
            table_context=table_ctx,
            column_mappings=column_mappings,
            tolerance=tolerance,
        )

        self.logger.info(
            f"SDV metrics validation for {table_name}: "
            f"is_valid={validation.is_valid}, rows={len(validation.results_dataframe)}"
        )
        return adapt_sdv_result(validation.results_dataframe, table_name)

    # ------------------------------------------------------------------
    # Query execution
    # ------------------------------------------------------------------

    def _run_source_query(self, query: str, task: dict[str, Any]) -> pd.DataFrame:
        """Run a query against the source database via ODBC."""
        if not query:
            return pd.DataFrame()

        payload = task.get(task_keys.PAYLOAD, task)
        engine = payload.get(PAYLOAD_SOURCE_PLATFORM, "")
        use_database_command = payload.get(PAYLOAD_USE_DATABASE_COMMAND)

        if engine not in self.source_connections_config:
            raise ValueError(
                f"Engine '{engine}' not found in source connections. "
                f"Available: {list(self.source_connections_config.keys())}"
            )

        import pyodbc

        connection_string = self.source_connections_config[engine].connection_string
        self.logger.info(f"Opening ODBC connection for engine '{engine}'")

        conn = pyodbc.connect(connection_string)
        register_output_converters(conn)

        try:
            if use_database_command:
                cursor = conn.cursor()
                cursor.execute(use_database_command)
                cursor.close()

            cursor = conn.cursor()
            self.logger.info(f"Executing source query: {query[:200]}...")
            cursor.execute(query)
            columns = [desc[0].upper() for desc in cursor.description]
            rows = cursor.fetchall()
            cursor.close()

            return pd.DataFrame.from_records(rows, columns=columns)
        finally:
            conn.close()

    def _run_target_query(self, query: str) -> pd.DataFrame:
        """Run a query against the target Snowflake database."""
        if not query:
            return pd.DataFrame()
        with self.snowflake_datasource.get_cursor(dict_cursor=False) as cursor:
            self.logger.info(f"Executing target query: {query[:200]}...")
            cursor.execute(query)
            columns = [desc[0] for desc in cursor.description]
            rows = cursor.fetchall()
            self.logger.info(
                f"Target query returned {len(rows)} rows, columns={columns}"
            )
            return pd.DataFrame(rows, columns=columns)

    def _open_source_connection(self, task: dict[str, Any]) -> Any:
        """Open and return a persistent connection to the source database."""
        import pyodbc

        payload = task.get(task_keys.PAYLOAD, task)
        engine = payload.get(PAYLOAD_SOURCE_PLATFORM, "")
        use_database_command = payload.get(PAYLOAD_USE_DATABASE_COMMAND)

        connection_string = self.source_connections_config[engine].connection_string
        conn = pyodbc.connect(connection_string)
        register_output_converters(conn)

        if use_database_command:
            cursor = conn.cursor()
            cursor.execute(use_database_command)
            cursor.close()

        return conn

    def _execute_source_statement_on_conn(
        self, conn: "pyodbc.Connection", statement: str
    ) -> None:
        """Execute a DDL/DML statement against the source using an existing connection."""
        cursor = conn.cursor()
        self.logger.info(f"Executing source statement: {statement[:200]}...")
        cursor.execute(statement)
        cursor.close()

    def _run_source_query_on_conn(
        self, conn: "pyodbc.Connection", query: str
    ) -> pd.DataFrame:
        """Run a query against the source using an existing connection."""
        cursor = conn.cursor()
        self.logger.info(f"Executing source query: {query[:200]}...")
        cursor.execute(query)
        columns = [desc[0].upper() for desc in cursor.description]
        rows = cursor.fetchall()
        cursor.close()
        return pd.DataFrame.from_records(rows, columns=columns)

    def _execute_target_statement(self, statement: str) -> None:
        """Execute a DDL/DML statement against the Snowflake target."""
        with self.snowflake_datasource.get_cursor(dict_cursor=False) as cursor:
            self.logger.info(f"Executing target statement: {statement[:200]}...")
            cursor.execute(statement)

    def _get_row_count_on_conn(self, conn: "pyodbc.Connection", table_fqn: str) -> int:
        """Get the row count of a source table using an existing ODBC connection."""
        query = f"SELECT COUNT(*) FROM {table_fqn}"
        cursor = conn.cursor()
        cursor.execute(query)
        row = cursor.fetchone()
        cursor.close()
        return int(row[0]) if row else 0

    def _get_target_row_count(self, table_fqn: str) -> int:
        """Get the row count of a target Snowflake table."""
        query = f"SELECT COUNT(*) FROM {table_fqn}"
        with self.snowflake_datasource.get_cursor(dict_cursor=False) as cursor:
            cursor.execute(query)
            row = cursor.fetchone()
            return int(row[0]) if row else 0

    # ------------------------------------------------------------------
    # Result output
    # ------------------------------------------------------------------

    def _write_result_csv(
        self,
        result_df: pd.DataFrame,
        workflow_id: int,
        table_metadata_id: int,
    ) -> str:
        """Write results to CSV with orchestration columns prepended."""
        temp_dir = tempfile.mkdtemp()
        csv_path = os.path.join(temp_dir, "validation_result.csv")

        with open(csv_path, "w", newline="", encoding="utf-8") as f:
            writer = csv.writer(f, delimiter=CSV_DELIMITER)
            for _, row in result_df.iterrows():
                csv_row = [workflow_id, table_metadata_id, *row.values]
                writer.writerow(csv_row)

        return csv_path

    def _write_l3_csv(
        self,
        temp_dir: str,
        filename: str,
        data_rows: list[list],
        workflow_id: int,
        table_metadata_id: int,
    ) -> str:
        """Write an L3 CSV file with workflow/table_metadata_id prepended to each row."""
        csv_path = os.path.join(temp_dir, filename)
        with open(csv_path, "w", newline="", encoding="utf-8") as f:
            writer = csv.writer(f, delimiter=CSV_DELIMITER)
            for row in data_rows:
                writer.writerow([workflow_id, table_metadata_id, *row])
        return csv_path

    @staticmethod
    def _cleanup_temp_dir(temp_dir: str) -> None:
        """Remove a temporary directory and all its contents."""
        shutil.rmtree(temp_dir, ignore_errors=True)

    def _upload_to_stage(self, local_path: str, stage_path: str) -> None:
        """Upload a local file to a Snowflake stage."""
        with self.snowflake_datasource.get_cursor(dict_cursor=False) as cursor:
            put_command = f"PUT 'file://{local_path}' '{stage_path}/' AUTO_COMPRESS=TRUE OVERWRITE=TRUE"
            cursor.execute(put_command)
            self.logger.info(f"Uploaded validation results to {stage_path}")
