"""
Pure helper functions, constants, and type definitions for data validation.

This module contains no I/O or side-effects; every public symbol is either
a constant, a type alias, or a stateless function.
"""

from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, TypedDict

import pandas as pd

from snowflake.snowflake_data_validation.public import (
    ColumnMetadata,
    ObjectType,
    Origin,
    Platform,
    SQLQueriesTemplateGenerator,
    TableContext,
    TemplatesLoaderManager,
)
from snowflake.snowflake_data_validation.utils.helpers.helper_templates import (
    HelperTemplates,
)


if TYPE_CHECKING:
    pass

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

VALIDATION_TYPE_SCHEMA = "schema"
VALIDATION_TYPE_METRICS = "metrics"
VALIDATION_TYPE_ROW = "row"

CSV_DELIMITER = "\x1f"

# Payload field keys (match DataValidationTaskPayload field names)
PAYLOAD_VALIDATION_TYPE = "validation_type"
PAYLOAD_TABLE_NAME = "table_name"
PAYLOAD_WORKFLOW_ID = "workflow_id"
PAYLOAD_TABLE_METADATA_ID = "table_metadata_id"
PAYLOAD_SOURCE_QUERY = "source_query"
PAYLOAD_TARGET_QUERY = "target_query"
PAYLOAD_TARGET_ID = "target_id"
PAYLOAD_SOURCE_PLATFORM = "source_platform"
PAYLOAD_COLUMN_MAPPINGS = "column_mappings"
PAYLOAD_DATATYPES_MAPPINGS = "datatypes_mappings"
PAYLOAD_TOLERANCE = "tolerance"
PAYLOAD_ROW_VALIDATION_MODE = "row_validation_mode"
PAYLOAD_INDEX_COLUMN_LIST = "index_column_list"
PAYLOAD_TARGET_INDEX_COLUMN_LIST = "target_index_column_list"
PAYLOAD_SOURCE_TABLE_FQN = "source_table_fqn"
PAYLOAD_TARGET_TABLE_FQN = "target_table_fqn"
PAYLOAD_CHUNK_NUMBER = "chunk_number"
PAYLOAD_MAX_FAILED_ROWS_NUMBER = "max_failed_rows_number"
PAYLOAD_COLUMNS_METADATA = "columns_metadata"
PAYLOAD_SOURCE_WHERE_CLAUSE = "source_where_clause"
PAYLOAD_TARGET_WHERE_CLAUSE = "target_where_clause"
PAYLOAD_USE_DATABASE_COMMAND = "use_database_command"
PAYLOAD_PARTITION_NUMBER = "partition_number"

_PLATFORM_MAP: dict[str, Platform] = {
    "redshift": Platform.REDSHIFT,
    "snowflake": Platform.SNOWFLAKE,
    "sqlserver": Platform.SQLSERVER,
    "sql_server": Platform.SQLSERVER,
    "teradata": Platform.TERADATA,
}

_PLATFORM_DIRS: dict[Platform, str] = {
    Platform.REDSHIFT: "redshift",
    Platform.SNOWFLAKE: "snowflake",
    Platform.SQLSERVER: "sqlserver",
    Platform.TERADATA: "teradata",
}

_DATATYPES_MAPPING_FORMAT = "{source_platform}_to_{target_platform}_datatypes_mapping_template.yaml"

# ---------------------------------------------------------------------------
# Type definitions
# ---------------------------------------------------------------------------


class TaskResult(TypedDict, total=False):
    """Return type for the data validation handler."""

    status: str
    details: str
    total: int


class ColumnMetadataDict(TypedDict):
    """Shape of each entry in the ``columns_metadata`` payload list."""

    name: str
    data_type: str


@dataclass
class MD5PipelineParams:
    """Groups the parameters for a single MD5 hash comparison pipeline run."""

    source_conn: Any  # pyodbc.Connection (import only at TYPE_CHECKING)
    source_ctx: TableContext
    target_ctx: TableContext
    index_column_list: list[str]
    target_index_column_list: list[str]
    source_table_fqn: str
    target_table_fqn: str
    table_name: str
    workflow_id: int
    table_metadata_id: int
    target_id: str
    max_failed_rows: int


# ---------------------------------------------------------------------------
# SDV integration helpers
# ---------------------------------------------------------------------------


def _sdv_package_root() -> Path:
    import snowflake.snowflake_data_validation as _sdv

    return Path(_sdv.__file__).parent


def _templates_path(platform: Platform) -> str:
    subdir = _PLATFORM_DIRS.get(platform, platform.value)
    return str(_sdv_package_root() / subdir / "extractor" / "templates")


def resolve_platform(platform_str: str) -> Platform:
    """Map a platform string to the corresponding SDV ``Platform`` enum."""
    p = _PLATFORM_MAP.get(platform_str.lower())
    if p is None:
        raise ValueError(f"Unknown platform: {platform_str}")
    return p


def build_table_context(
    platform: Platform,
    table_name: str,
) -> TableContext:
    """Minimal TableContext sufficient for validate_schema / validate_metrics calls."""
    templates_dir = _templates_path(platform)
    sql_generator = SQLQueriesTemplateGenerator(jinja_templates_folder_path=Path(templates_dir))
    templates_loader_manager = TemplatesLoaderManager(
        templates_directory_path=Path(templates_dir),
        platform=platform,
    )
    return TableContext(
        apply_metric_column_modifier=True,
        chunk_number=1,
        column_mappings={},
        column_selection_list=[],
        columns=[],
        database_name="",
        exclude_metrics=False,
        fully_qualified_name=table_name,
        has_where_clause=False,
        id=0,
        internal_fully_qualified_name=table_name,
        internal_table_name=table_name,
        is_case_sensitive=False,
        is_exclusion_mode=False,
        is_temporary_table=False,
        max_failed_rows_number=100,
        object_type=ObjectType.TABLE,
        origin=Origin.SOURCE,
        platform=platform,
        row_count=0,
        run_id="0",
        run_start_time="",
        schema_name="",
        sql_generator=sql_generator,
        table_name=table_name,
        templates_loader_manager=templates_loader_manager,
        user_index_column_collection=[],
        where_clause="",
    )


def build_row_table_context(
    platform: Platform,
    origin: Origin,
    table_fqn: str,
    columns: list[ColumnMetadata],
    index_column_list: list[str],
    chunk_number: int,
    max_failed_rows_number: int,
    row_count: int = 0,
    where_clause: str = "",
) -> TableContext:
    """Build a TableContext with column/index metadata needed by the MD5 pipeline."""
    templates_dir = _templates_path(platform)
    sql_generator = SQLQueriesTemplateGenerator(jinja_templates_folder_path=Path(templates_dir))
    templates_loader_manager = TemplatesLoaderManager(
        templates_directory_path=Path(templates_dir),
        platform=platform,
    )

    parts = table_fqn.split(".")
    if len(parts) == 3:
        db_name, schema_name, table_name = parts
    elif len(parts) == 2:
        db_name, schema_name, table_name = "", parts[0], parts[1]
    else:
        db_name, schema_name, table_name = "", "", table_fqn

    return TableContext(
        apply_metric_column_modifier=True,
        chunk_number=chunk_number,
        column_mappings={},
        column_selection_list=[],
        columns=columns,
        database_name=db_name,
        exclude_metrics=False,
        fully_qualified_name=table_fqn,
        has_where_clause=bool(where_clause),
        id=0,
        internal_fully_qualified_name=table_fqn,
        internal_table_name=table_name,
        is_case_sensitive=False,
        is_exclusion_mode=False,
        is_temporary_table=False,
        max_failed_rows_number=max_failed_rows_number,
        object_type=ObjectType.TABLE,
        origin=origin,
        platform=platform,
        row_count=row_count,
        run_id="0",
        run_start_time="",
        schema_name=schema_name,
        sql_generator=sql_generator,
        table_name=table_name,
        templates_loader_manager=templates_loader_manager,
        user_index_column_collection=index_column_list,
        where_clause=where_clause,
    )


def load_datatypes_mappings(source_platform: Platform) -> dict[str, str]:
    """Load the YAML datatype mapping for *source_platform* -> Snowflake."""
    file_name = _DATATYPES_MAPPING_FORMAT.format(
        source_platform=source_platform.value,
        target_platform=Platform.SNOWFLAKE.value,
    )
    mapping_path = Path(_templates_path(source_platform)) / file_name
    if not mapping_path.exists():
        return {}
    return HelperTemplates.load_datatypes_mapping_templates_from_yaml(yaml_path=mapping_path)


def adapt_sdv_result(
    result_df: pd.DataFrame,
    table_name: str,
) -> pd.DataFrame:
    """
    Rename / fill columns so the DataFrame matches the COPY INTO schema.

    SDV output uses ``TABLE`` for the table-name column; our Snowflake tables
    expect ``TABLE_NAME``.  This function also ensures all expected columns
    are present.
    """
    if result_df.empty:
        return pd.DataFrame(
            columns=[
                "VALIDATION_TYPE",
                "TABLE_NAME",
                "COLUMN_VALIDATED",
                "EVALUATION_CRITERIA",
                "SOURCE_VALUE",
                "SNOWFLAKE_VALUE",
                "STATUS",
                "COMMENTS",
            ]
        )

    if "TABLE" in result_df.columns and "TABLE_NAME" not in result_df.columns:
        result_df = result_df.rename(columns={"TABLE": "TABLE_NAME"})

    if "TABLE_NAME" not in result_df.columns:
        result_df["TABLE_NAME"] = table_name

    expected = [
        "VALIDATION_TYPE",
        "TABLE_NAME",
        "COLUMN_VALIDATED",
        "EVALUATION_CRITERIA",
        "SOURCE_VALUE",
        "SNOWFLAKE_VALUE",
        "STATUS",
        "COMMENTS",
    ]
    for col in expected:
        if col not in result_df.columns:
            result_df[col] = ""

    return result_df[expected]
