"""
L3 cell-by-cell validation handler.

Compares individual cell values between source and target using the Polars-based
``validate_cell`` function from the SDV library.
"""

from __future__ import annotations

import tempfile

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

import pandas as pd
import polars as pl

from data_exchange_agent.constants import task_keys
from data_exchange_agent.tasks.data_validation.helpers import (
    PAYLOAD_SOURCE_PLATFORM,
    PAYLOAD_SOURCE_QUERY,
    PAYLOAD_TABLE_METADATA_ID,
    PAYLOAD_TABLE_NAME,
    PAYLOAD_TARGET_ID,
    PAYLOAD_TARGET_QUERY,
    PAYLOAD_WORKFLOW_ID,
    TaskResult,
    build_table_context,
    resolve_platform,
)
from snowflake.snowflake_data_validation.public import (
    QueryExecutionResult,
    validate_cell,
)


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


class CellValidationHandler:
    """Encapsulates L3 cell-by-cell 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

    def handle(self, task: dict[str, Any], payload: dict[str, Any]) -> TaskResult:
        """Execute cell-by-cell validation using Polars-based comparison."""
        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, "")
        source_platform_str = payload.get(PAYLOAD_SOURCE_PLATFORM, "")

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

        self.logger.info(f"Starting cell validation for {table_name}")

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

        target_df.columns = [c.upper() for c in target_df.columns]

        self.logger.info(f"Cell validation DataFrames: source={source_df.shape}, target={target_df.shape}")

        source_pl = pl.from_pandas(source_df)
        target_pl = pl.from_pandas(target_df)

        src_result = QueryExecutionResult(dataframe=source_pl)
        tgt_result = QueryExecutionResult(dataframe=target_pl)

        validation = validate_cell(src_result, tgt_result, table_ctx)

        result_df = validation.results_dataframe
        self.logger.info(f"Cell validation for {table_name}: " f"is_valid={validation.is_valid}, diffs={len(result_df)}")

        diffs_data: list[list] = []
        if not result_df.empty:
            diffs_df = pd.DataFrame(
                {
                    "table": table_name,
                    "row_num": result_df["ROW"].fillna(0),
                    "column": result_df["COLUMN"].fillna(""),
                    "source_value": result_df["SOURCE_VALUE"].fillna("").astype(str),
                    "target_value": result_df["TARGET_VALUE"].fillna("").astype(str),
                }
            )
            diffs_data = diffs_df.values.tolist()

        temp_dir = tempfile.mkdtemp()
        diffs_path = self._parent._write_l3_csv(temp_dir, "diffs.csv", diffs_data, workflow_id, table_metadata_id)
        self._parent._upload_to_stage(diffs_path, f"{target_id}/diffs")
        self._parent._cleanup_temp_dir(temp_dir)

        self.logger.info(f"Cell validation complete for {table_name}: {len(result_df)} diffs uploaded")
        return {
            task_keys.STATUS: "success",
            task_keys.DETAILS: f"Cell validation: {len(result_df)} diffs for {table_name}",
            task_keys.TOTAL: len(result_df),
        }
