"""
L3 row-by-row (MD5) validation handler.

Implements the full MD5 chunk pipeline: create staging tables, compute hashes,
compare chunk-level hashes, drill down into differing chunks, and upload
diff/summary/hashes CSVs to the Snowflake stage.
"""

from __future__ import annotations

import tempfile

from typing import TYPE_CHECKING, Any

import pandas as pd

from data_exchange_agent.constants import task_keys
from data_exchange_agent.tasks.data_validation.helpers import (
    PAYLOAD_CHUNK_NUMBER,
    PAYLOAD_COLUMNS_METADATA,
    PAYLOAD_INDEX_COLUMN_LIST,
    PAYLOAD_MAX_FAILED_ROWS_NUMBER,
    PAYLOAD_SOURCE_PLATFORM,
    PAYLOAD_SOURCE_TABLE_FQN,
    PAYLOAD_SOURCE_WHERE_CLAUSE,
    PAYLOAD_TABLE_METADATA_ID,
    PAYLOAD_TABLE_NAME,
    PAYLOAD_TARGET_ID,
    PAYLOAD_TARGET_INDEX_COLUMN_LIST,
    PAYLOAD_TARGET_TABLE_FQN,
    PAYLOAD_TARGET_WHERE_CLAUSE,
    PAYLOAD_WORKFLOW_ID,
    MD5PipelineParams,
    TaskResult,
    build_row_table_context,
    resolve_platform,
)
from snowflake.snowflake_data_validation.public import (
    ColumnMetadata,
    Origin,
    Platform,
    QueryExecutionResult,
    compare_md5_chunks,
    compare_md5_rows_chunk,
    generate_compute_md5_queries,
    generate_create_md5_table_statement,
    generate_extract_chunks_md5_query,
    generate_extract_md5_rows_chunk_query,
)
from snowflake.snowflake_data_validation.utils.constants import (
    CHUNK_ID_COLUMN_KEY,
    CHUNK_MD5_VALUE_COLUMN_KEY,
)


if TYPE_CHECKING:
    from data_exchange_agent.tasks.data_validation.handler import (
        DataValidationHandler,
    )


class RowValidationHandler:
    """Encapsulates L3 row-by-row MD5 validation logic."""

    def __init__(self, parent: DataValidationHandler) -> None:
        """Initialize with a reference to the parent handler."""
        self._parent = parent

    @property
    def logger(self):
        """Return the parent handler's logger."""
        return self._parent.logger

    # ------------------------------------------------------------------
    # Public entry point
    # ------------------------------------------------------------------

    def handle(self, task: dict[str, Any], payload: dict[str, Any]) -> TaskResult:
        """Execute the full MD5 chunk pipeline for row-by-row validation."""
        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)
        target_id = payload.get(PAYLOAD_TARGET_ID, "")
        source_platform_str = payload.get(PAYLOAD_SOURCE_PLATFORM, "")
        source_table_fqn = payload.get(PAYLOAD_SOURCE_TABLE_FQN, table_name)
        target_table_fqn = payload.get(PAYLOAD_TARGET_TABLE_FQN, table_name)
        index_column_list = payload.get(PAYLOAD_INDEX_COLUMN_LIST) or []
        target_index_column_list = payload.get(PAYLOAD_TARGET_INDEX_COLUMN_LIST) or index_column_list
        chunk_number = payload.get(PAYLOAD_CHUNK_NUMBER) or 1
        max_failed_rows = payload.get(PAYLOAD_MAX_FAILED_ROWS_NUMBER) or 100
        columns_metadata_raw = payload.get(PAYLOAD_COLUMNS_METADATA) or []
        source_where_clause = payload.get(PAYLOAD_SOURCE_WHERE_CLAUSE) or ""
        target_where_clause = payload.get(PAYLOAD_TARGET_WHERE_CLAUSE) or ""

        if not index_column_list:
            self.logger.error(f"Row validation requires index_column_list for {table_name}")
            return {
                task_keys.STATUS: "error",
                task_keys.DETAILS: "Missing index_column_list",
            }

        source_platform = resolve_platform(source_platform_str)

        source_columns = [
            ColumnMetadata(
                name=c["name"],
                data_type=c["data_type"],
                nullable=True,
                is_primary_key=False,
                calculated_column_size_in_bytes=0,
                properties={},
            )
            for c in columns_metadata_raw
        ]

        target_columns = self._get_target_column_metadata(target_table_fqn, source_columns)

        self.logger.info(f"Starting MD5 row validation for {source_table_fqn} -> {target_table_fqn}")

        source_conn = self._parent._open_source_connection(task)

        source_row_count = self._parent._get_row_count_on_conn(source_conn, source_table_fqn)
        target_row_count = self._parent._get_target_row_count(target_table_fqn)
        self.logger.info(f"Row counts for MD5 pipeline: source={source_row_count}, target={target_row_count}")

        source_ctx = build_row_table_context(
            platform=source_platform,
            origin=Origin.SOURCE,
            table_fqn=source_table_fqn,
            columns=source_columns,
            index_column_list=index_column_list,
            chunk_number=chunk_number,
            max_failed_rows_number=max_failed_rows,
            row_count=source_row_count,
            where_clause=source_where_clause,
        )
        target_ctx = build_row_table_context(
            platform=Platform.SNOWFLAKE,
            origin=Origin.TARGET,
            table_fqn=target_table_fqn,
            columns=target_columns,
            index_column_list=target_index_column_list,
            chunk_number=chunk_number,
            max_failed_rows_number=max_failed_rows,
            row_count=target_row_count,
            where_clause=target_where_clause,
        )
        pipeline_params = MD5PipelineParams(
            source_conn=source_conn,
            source_ctx=source_ctx,
            target_ctx=target_ctx,
            index_column_list=index_column_list,
            target_index_column_list=target_index_column_list,
            source_table_fqn=source_table_fqn,
            target_table_fqn=target_table_fqn,
            table_name=table_name,
            workflow_id=workflow_id,
            table_metadata_id=table_metadata_id,
            target_id=target_id,
            max_failed_rows=max_failed_rows,
        )
        try:
            return self._run_md5_pipeline(pipeline_params)
        finally:
            source_conn.close()

    # ------------------------------------------------------------------
    # MD5 pipeline
    # ------------------------------------------------------------------

    def _run_md5_pipeline(self, params: MD5PipelineParams) -> TaskResult:
        """Execute the full MD5 pipeline using a persistent source connection."""
        source_conn = params.source_conn
        source_ctx = params.source_ctx
        target_ctx = params.target_ctx
        index_column_list = params.index_column_list
        target_index_column_list = params.target_index_column_list
        source_table_fqn = params.source_table_fqn
        target_table_fqn = params.target_table_fqn
        table_name = params.table_name
        workflow_id = params.workflow_id
        table_metadata_id = params.table_metadata_id
        target_id = params.target_id
        max_failed_rows = params.max_failed_rows

        src_idx_cols = [col.upper() + "_SOURCE" for col in index_column_list]
        tgt_idx_cols = [col.upper() + "_TARGET" for col in target_index_column_list]

        # Step 1: Create MD5 staging tables on both sides
        src_create = generate_create_md5_table_statement(source_ctx)
        tgt_create = generate_create_md5_table_statement(target_ctx)
        self._parent._execute_source_statement_on_conn(source_conn, src_create.query_string)
        self._parent._execute_target_statement(tgt_create.query_string)

        # Step 2: Compute MD5 chunks on both sides
        src_compute_queries = generate_compute_md5_queries(source_ctx, target_ctx.table_name)
        tgt_compute_queries = generate_compute_md5_queries(target_ctx, source_ctx.table_name)

        for q in src_compute_queries:
            self._parent._execute_source_statement_on_conn(source_conn, q.query_string)
        for q in tgt_compute_queries:
            self._parent._execute_target_statement(q.query_string)

        # Step 3: Extract chunk MD5s
        src_extract = generate_extract_chunks_md5_query(source_ctx)
        tgt_extract = generate_extract_chunks_md5_query(target_ctx)

        source_chunks_df = self._parent._run_source_query_on_conn(source_conn, src_extract.query_string)
        target_chunks_df = self._parent._run_target_query(tgt_extract.query_string)

        if source_chunks_df.empty or target_chunks_df.empty:
            self.logger.warning(f"Empty chunk MD5 results for {table_name}, writing empty results")
            self._write_empty_row_results(workflow_id, table_metadata_id, table_name, target_table_fqn, target_id)
            return {
                task_keys.STATUS: "success",
                task_keys.DETAILS: f"Row validation: no chunk data for {table_name}",
            }

        # Step 4: Compare chunks
        src_chunks_result = QueryExecutionResult(dataframe=source_chunks_df)
        tgt_chunks_result = QueryExecutionResult(dataframe=target_chunks_df)
        chunk_comparison = compare_md5_chunks(src_chunks_result, tgt_chunks_result, source_ctx)

        hashes_df = self.build_hashes_dataframe(source_chunks_df, target_chunks_df, table_name)

        # Step 5: Drill down into differing chunks
        all_diffs: list[pd.DataFrame] = []
        total_rows_compared = 0
        rows_not_found_in_target = 0
        rows_not_found_in_source = 0
        rows_with_differences = 0

        if not chunk_comparison.is_valid:
            diff_chunks_df = chunk_comparison.results_dataframe

            for _, chunk_row in diff_chunks_df.iterrows():
                chunk_id = chunk_row.get(
                    CHUNK_ID_COLUMN_KEY,
                    chunk_row.get(f"{CHUNK_ID_COLUMN_KEY}_SOURCE", ""),
                )
                if not chunk_id:
                    continue

                self.logger.info(f"Drilling into chunk {chunk_id}")
                src_rows_query = generate_extract_md5_rows_chunk_query(chunk_id, source_ctx)
                tgt_rows_query = generate_extract_md5_rows_chunk_query(chunk_id, target_ctx)

                src_rows_df = self._parent._run_source_query_on_conn(source_conn, src_rows_query.query_string)
                tgt_rows_df = self._parent._run_target_query(tgt_rows_query.query_string)

                src_rows_df = src_rows_df.convert_dtypes().rename(columns=lambda x: x.upper() + "_SOURCE")
                tgt_rows_df = tgt_rows_df.convert_dtypes().rename(columns=lambda x: x.upper() + "_TARGET")

                src_rows_result = QueryExecutionResult(dataframe=src_rows_df)
                tgt_rows_result = QueryExecutionResult(dataframe=tgt_rows_df)

                row_diff = compare_md5_rows_chunk(
                    src_rows_result,
                    tgt_rows_result,
                    src_idx_cols,
                    tgt_idx_cols,
                    source_ctx,
                )

                if not row_diff.results_dataframe.empty:
                    diff_df = row_diff.results_dataframe
                    total_rows_compared += len(src_rows_df) + len(tgt_rows_df)
                    counts = diff_df["RESULT"].value_counts()
                    rows_not_found_in_target += int(counts.get("NOT_FOUND_TARGET", 0))
                    rows_not_found_in_source += int(counts.get("NOT_FOUND_SOURCE", 0))
                    rows_with_differences += int(counts.get("FAILURE", 0))
                    all_diffs.append(diff_df)

                    if sum(len(d) for d in all_diffs) >= max_failed_rows:
                        self.logger.info(f"Reached max_failed_rows ({max_failed_rows}), stopping drill-down")
                        break

        # Step 6: Build diffs CSV
        diffs_data: list[list] = []
        if all_diffs:
            combined = pd.concat(all_diffs, ignore_index=True)
            src_parts = [combined[col].fillna("N/A").astype(str).radd(f"{col.replace('_SOURCE', '')}=") for col in src_idx_cols]
            tgt_parts = [combined[col].fillna("N/A").astype(str).radd(f"{col.replace('_TARGET', '')}=") for col in tgt_idx_cols]
            src_vals_series = src_parts[0].str.cat(src_parts[1:], sep=", ") if src_parts else pd.Series("", index=combined.index)
            tgt_vals_series = tgt_parts[0].str.cat(tgt_parts[1:], sep=", ") if tgt_parts else pd.Series("", index=combined.index)

            diffs_out = pd.DataFrame(
                {
                    "source_fqn": source_table_fqn,
                    "object_type": "TABLE",
                    "result": combined["RESULT"].fillna(""),
                    "src_vals": src_vals_series,
                    "tgt_vals": tgt_vals_series,
                    "target_fqn": target_table_fqn,
                }
            )
            diffs_data = diffs_out.values.tolist()

        total_chunks = max(len(source_chunks_df), len(target_chunks_df))
        differing = total_chunks - int(hashes_df["IS_MATCH"].sum()) if not hashes_df.empty else 0
        matching = total_chunks - differing

        summary_data = [
            [
                table_name,
                total_chunks,
                matching,
                differing,
                total_rows_compared,
                rows_not_found_in_target,
                rows_not_found_in_source,
                rows_with_differences,
                chunk_comparison.is_valid,
            ]
        ]

        # Step 7: Write and upload CSVs
        temp_dir = tempfile.mkdtemp()

        diffs_path = self._parent._write_l3_csv(temp_dir, "diffs.csv", diffs_data, workflow_id, table_metadata_id)
        summary_path = self._parent._write_l3_csv(temp_dir, "summary.csv", summary_data, workflow_id, table_metadata_id)
        hashes_path = self._parent._write_l3_csv(
            temp_dir,
            "hashes.csv",
            hashes_df.values.tolist(),
            workflow_id,
            table_metadata_id,
        )

        self._parent._upload_to_stage(diffs_path, f"{target_id}/diffs")
        self._parent._upload_to_stage(summary_path, f"{target_id}/summary")
        self._parent._upload_to_stage(hashes_path, f"{target_id}/hashes")

        self._parent._cleanup_temp_dir(temp_dir)

        total_diffs = sum(len(d) for d in all_diffs)
        self.logger.info(
            f"Row validation complete for {table_name}: " f"{total_chunks} chunks, {differing} differing, {total_diffs} row diffs"
        )
        return {
            task_keys.STATUS: "success",
            task_keys.DETAILS: f"Row validation: {total_diffs} diffs in {differing}/{total_chunks} chunks",
            task_keys.TOTAL: total_diffs,
        }

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

    @staticmethod
    def build_hashes_dataframe(
        source_chunks_df: pd.DataFrame,
        target_chunks_df: pd.DataFrame,
        table_name: str,
    ) -> pd.DataFrame:
        """Build a chunk hash comparison DataFrame from both sides' extracted MD5 data."""
        src_map = dict(
            zip(
                source_chunks_df[CHUNK_ID_COLUMN_KEY],
                source_chunks_df[CHUNK_MD5_VALUE_COLUMN_KEY],
                strict=False,
            )
        )
        tgt_map = dict(
            zip(
                target_chunks_df[CHUNK_ID_COLUMN_KEY],
                target_chunks_df[CHUNK_MD5_VALUE_COLUMN_KEY],
                strict=False,
            )
        )

        all_chunk_ids = sorted(set(src_map) | set(tgt_map))
        rows = [
            [
                table_name,
                cid,
                src_map.get(cid, ""),
                tgt_map.get(cid, ""),
                src_map.get(cid, "") == tgt_map.get(cid, "") and src_map.get(cid, "") != "",
            ]
            for cid in all_chunk_ids
        ]
        return pd.DataFrame(
            rows,
            columns=[
                "TABLE_NAME",
                "CHUNK_ID",
                "SOURCE_HASH",
                "TARGET_HASH",
                "IS_MATCH",
            ],
        )

    def _write_empty_row_results(
        self,
        workflow_id: int,
        table_metadata_id: int,
        table_name: str,
        target_table_fqn: str,
        target_id: str,
    ) -> None:
        """Write empty CSV files for row validation when no data is available."""
        temp_dir = tempfile.mkdtemp()
        empty_diffs = self._parent._write_l3_csv(temp_dir, "diffs.csv", [], workflow_id, table_metadata_id)
        empty_summary = self._parent._write_l3_csv(
            temp_dir,
            "summary.csv",
            [[table_name, 0, 0, 0, 0, 0, 0, 0, True]],
            workflow_id,
            table_metadata_id,
        )
        empty_hashes = self._parent._write_l3_csv(temp_dir, "hashes.csv", [], workflow_id, table_metadata_id)

        self._parent._upload_to_stage(empty_diffs, f"{target_id}/diffs")
        self._parent._upload_to_stage(empty_summary, f"{target_id}/summary")
        self._parent._upload_to_stage(empty_hashes, f"{target_id}/hashes")
        self._parent._cleanup_temp_dir(temp_dir)

    def _get_target_column_metadata(
        self,
        target_table_fqn: str,
        fallback_columns: list[ColumnMetadata],
    ) -> list[ColumnMetadata]:
        """Query Snowflake information_schema for the target table's native column types."""
        parts = target_table_fqn.split(".")
        if len(parts) == 3:
            tgt_db, tgt_schema, tgt_table = parts
        elif len(parts) == 2:
            tgt_db, tgt_schema, tgt_table = "", parts[0], parts[1]
        else:
            return fallback_columns

        prefix = f"{tgt_db}." if tgt_db else ""
        query = (
            f"SELECT column_name, data_type "
            f"FROM {prefix}information_schema.columns "
            f"WHERE table_schema = '{tgt_schema.upper()}' "
            f"AND table_name = '{tgt_table.upper()}' "
            f"ORDER BY ordinal_position"
        )
        try:
            df = self._parent._run_target_query(query)
            if df.empty:
                self.logger.warning(
                    f"No columns found in information_schema for {target_table_fqn}, " "falling back to source column types"
                )
                return fallback_columns

            columns = []
            for _, row in df.iterrows():
                col_name = row.get("COLUMN_NAME", row.get("column_name", ""))
                data_type = row.get("DATA_TYPE", row.get("data_type", ""))
                if col_name and data_type:
                    columns.append(
                        ColumnMetadata(
                            name=col_name,
                            data_type=data_type,
                            nullable=True,
                            is_primary_key=False,
                            calculated_column_size_in_bytes=0,
                            properties={},
                        )
                    )
            return columns if columns else fallback_columns
        except Exception as e:
            self.logger.warning(f"Failed to get target column metadata for {target_table_fqn}: {e}, " "falling back to source column types")
            return fallback_columns
