"""
Transient table management for Snowflake-native validation.

Creates and cleans up transient tables for both baseline and actual
result sets so the data-validation framework's ``MetadataExtractorSnowflake``
can query them via SQL.

Baseline tables are built directly from the staged JSON files using
``LATERAL FLATTEN``, avoiding any intermediate table.

When capture or fetch statements are present, the **last** result set
(and its column types) is used for comparison.  For legacy capture
baselines, ``result_sets[0]`` is the OUT row and ``result_sets[-1]`` is
the data.  For template-driven baselines (``output.tables`` /
``output.cursors``), scalar params are in ``output_params`` and each
``result_sets`` entry corresponds to one table or cursor FETCH.
"""

from __future__ import annotations

import json
import logging
import time
from datetime import datetime as _datetime

from snowflake.connector.cursor import SnowflakeCursor
from snowflake.snowflake_data_validation.connector.connector_base import ConnectorBase


LOGGER = logging.getLogger(__name__)

_SCHEMA = "VALIDATION"

# Snowflake SQL type families used to classify columns for stats and coercion.
_NUMERIC_TYPE_PREFIXES = (
    "NUMBER", "NUMERIC", "DECIMAL",
    "INT", "INTEGER", "BIGINT", "SMALLINT", "TINYINT", "BYTEINT",
    "FLOAT", "FLOAT4", "FLOAT8", "DOUBLE", "REAL",
)

_TIMESTAMP_TYPE_PREFIXES = ("TIMESTAMP",)


def is_numeric_type(sql_type: str) -> bool:
    """Return True if *sql_type* is a numeric Snowflake data type."""
    base = sql_type.upper().split("(")[0].strip()
    return base in _NUMERIC_TYPE_PREFIXES or base.startswith(_NUMERIC_TYPE_PREFIXES)


def is_timestamp_type(sql_type: str) -> bool:
    """Return True if *sql_type* is a Snowflake timestamp data type."""
    base = sql_type.upper().split("(")[0].strip()
    return base.startswith(_TIMESTAMP_TYPE_PREFIXES)


def coerce_timestamp(value: object) -> object:
    """Parse ISO 8601 timestamp strings to Python datetime objects.

    Baseline JSON stores datetimes as ISO strings (e.g. ``'2026-01-15T10:30:00'``).
    Snowflake returns ``datetime`` objects for TIMESTAMP columns, which Polars
    infers as ``Datetime`` dtype.  Parsing baseline strings to ``datetime``
    ensures both sides get the same Polars column type, preventing false
    type-mismatch failures in ``validate_cell``.
    """
    if value is None or isinstance(value, _datetime):
        return value
    if isinstance(value, str):
        try:
            return _datetime.fromisoformat(value)
        except (ValueError, TypeError):
            pass
    return value


def coerce_numeric(value: object) -> object:
    """Cast numeric strings to ``int`` or ``float``; return everything else unchanged."""
    if value is None or isinstance(value, (int, float)):
        return value
    s = str(value).strip()
    try:
        # Prefer int for whole-number strings ("4", "42")
        return int(s)
    except ValueError:
        pass
    try:
        return float(s)
    except ValueError:
        return value


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

def _make_table_name(database: str, procedure: str, params_hash: str, side: str) -> str:
    """Generate a unique transient table name."""
    safe = procedure.replace(".", "_").replace(" ", "_")[:50].upper()
    return f"{database}.{_SCHEMA}.TEMP_{side}_{safe}_{params_hash}".upper()


def _safe_cast_type(typ: str) -> str:
    """Widen bare ``NUMBER`` to ``DOUBLE`` for VARIANT casts.

    Snowflake defaults bare ``NUMBER`` to ``NUMBER(38,0)`` which
    truncates decimals when casting from VARIANT/JSON.  Using ``DOUBLE``
    preserves the original decimal values stored in the baseline JSON
    without imposing a fixed scale that could mismatch the target table's
    native column types.
    """
    if typ.upper() == "NUMBER":
        return "DOUBLE"
    return typ


# ---------------------------------------------------------------------------
# Materialize baseline directly from stage JSON
# ---------------------------------------------------------------------------

def materialize_baseline_transient_table(
    connector: ConnectorBase,
    database: str,
    procedure: str,
    params_hash: str,
    target_col_types: dict[str, str] | None = None,
) -> str:
    """Create a transient table from baseline JSON files on the stage.

    Reads the staged baseline as raw text so ``json.loads`` preserves
    the original column order, then uses ``LATERAL FLATTEN`` with
    ``GET_IGNORE_CASE`` to materialize the transient table.

    When *target_col_types* is provided, baseline columns are cast to
    the target table's actual types so the validation framework sees
    matching column types on both sides.

    Args:
        connector: Snowflake connector.
        database: Database name for the validation schema.
        procedure: Procedure name to filter by.
        params_hash: Parameter hash to filter by.
        target_col_types: Column name -> Snowflake type mapping from the
            target transient table.  When provided, overrides the baseline
            JSON's stored types for casting.

    Returns:
        The fully qualified transient table name.
    """
    table_name = _make_table_name(database, procedure, params_hash, "SOURCE")
    fq = f"{database}.{_SCHEMA}"
    stage = f"@{fq}.BASELINES"
    fmt = f"{fq}.JSON_FORMAT"
    file_filter = f"s.$1:params_hash::VARCHAR = '{params_hash}'"

    raw_fmt = f"{fq}.RAW_TEXT_FORMAT"
    cursor = connector.connection.cursor()
    try:
        # Read the raw JSON text to preserve original key order
        # (Snowflake VARIANT alphabetizes object keys)
        cursor.execute(
            f"CREATE FILE FORMAT IF NOT EXISTS {raw_fmt} "
            f"TYPE = CSV FIELD_DELIMITER = NONE RECORD_DELIMITER = NONE"
        )
        cursor.execute(f"""\
SELECT $1 AS raw_text
FROM {stage} (FILE_FORMAT => '{raw_fmt}')
WHERE PARSE_JSON($1):params_hash::VARCHAR = '{params_hash}'
LIMIT 1""")
        row = cursor.fetchone()
        if row is None:
            LOGGER.warning("No baseline on stage for %s (hash=%s)", procedure, params_hash)
            cursor.execute(
                f"CREATE OR REPLACE TRANSIENT TABLE {table_name} (DUMMY VARCHAR)"
            )
            return table_name

        baseline: dict = json.loads(row[0])
        result_sets = baseline.get("result_sets", [])
        col_types_list: list[dict] = baseline.get("column_types", [{}])

        # Use the last result set for comparison.  For simple procedures
        # this is index 0 (the only set).  For legacy capture baselines,
        # index 0 is the OUT row and the last is the data.  For
        # template-driven baselines the last set is the final
        # table/cursor fetch.
        data_idx = len(result_sets) - 1 if result_sets else 0
        col_types: dict[str, str] = (
            col_types_list[data_idx] if data_idx < len(col_types_list) else {}
        )


        # When target column types are available, use them for casting
        # so both tables end up with matching types for the validator.
        cast_types = {}
        for col, typ in col_types.items():
            if target_col_types and col in target_col_types:
                cast_types[col] = target_col_types[col]
            else:
                cast_types[col] = _safe_cast_type(typ)

        select_cols = ", ".join(
            f"GET_IGNORE_CASE(f.value, '{col}')"
            f"::{cast_typ} "
            f'AS "{col}"'
            for col, cast_typ in cast_types.items()
        )

        cursor.execute(f"""\
CREATE OR REPLACE TRANSIENT TABLE {table_name} AS
SELECT {select_cols}
FROM {stage} (FILE_FORMAT => '{fmt}') s,
LATERAL FLATTEN(input => s.$1:result_sets[{data_idx}]) f
WHERE {file_filter}""")

    finally:
        cursor.close()

    return table_name


# ---------------------------------------------------------------------------
# Materialize actual results via CALL + RESULT_SCAN
# ---------------------------------------------------------------------------

def execute_with_capture(
    cursor: SnowflakeCursor,
    sql: str,
    fetch_stmts: list[str] | None = None,
) -> None:
    """Execute *sql* and optionally run follow-up statements on *cursor*.

    After return the cursor is positioned at the final result set — either
    the CALL's own result or the last follow-up statement's result.

    *fetch_stmts* are pre-rendered SQL statements (e.g. ``SELECT * FROM
    IDENTIFIER('...')``) produced by ``render_output_fetch_statements()``.
    """
    cursor.execute(sql)

    if fetch_stmts:
        for stmt in fetch_stmts:
            cursor.execute(stmt)


def materialize_actual_transient_table(
    connector: ConnectorBase,
    database: str,
    sql: str,
    procedure: str,
    params_hash: str,
    fetch_stmts: list[str] | None = None,
) -> str:
    """Execute a procedure and capture its results into a transient table.

    When *fetch_stmts* is provided, each pre-rendered statement (e.g.
    ``SELECT * FROM IDENTIFIER()``) is executed after the CALL and the
    **last** statement's result is materialized via ``RESULT_SCAN``.

    When *fetch_stmts* is not provided, the CALL's own result is materialized.

    Returns:
        The fully qualified transient table name.
    """
    table_name = _make_table_name(database, procedure, params_hash, "TARGET")

    cursor = connector.connection.cursor()
    try:
        execute_with_capture(cursor, sql, fetch_stmts)
        query_id = cursor.sfqid
        cursor.execute(
            f"CREATE OR REPLACE TRANSIENT TABLE {table_name} "
            f"AS SELECT * FROM TABLE(RESULT_SCAN('{query_id}'))"
        )
    finally:
        cursor.close()

    return table_name


# ---------------------------------------------------------------------------
# Fetch baseline rows into Python memory (for in-memory validation)
# ---------------------------------------------------------------------------

def fetch_baseline_rows(
    connector: ConnectorBase,
    database: str,
    params_hash: str,
) -> tuple[list[dict], dict[str, str]]:
    """Return ``(rows, col_types)`` for the baseline matching *params_hash*.

    Reads the stage JSON directly into Python — no Snowflake objects created.
    Used by the in-memory (DuckDB) comparison path.
    """
    fq = f"{database}.{_SCHEMA}"
    stage = f"@{fq}.BASELINES"
    raw_fmt = f"{fq}.RAW_TEXT_FORMAT"

    cursor = connector.connection.cursor()
    try:
        t0 = time.monotonic()
        # TYPE=CSV with no delimiters reads the whole file as a single raw string.
        # Snowflake's JSON file format parses into VARIANT which alphabetizes
        # object keys, breaking the column order stored in column_types.
        cursor.execute(
            f"CREATE FILE FORMAT IF NOT EXISTS {raw_fmt} "
            f"TYPE = CSV FIELD_DELIMITER = NONE RECORD_DELIMITER = NONE"
        )
        LOGGER.info("[timing] CREATE FILE FORMAT: %.3fs", time.monotonic() - t0)
        t1 = time.monotonic()
        cursor.execute(f"""\
SELECT $1 AS raw_text
FROM {stage} (FILE_FORMAT => '{raw_fmt}')
WHERE PARSE_JSON($1):params_hash::VARCHAR = '{params_hash}'
LIMIT 1""")
        LOGGER.info("[timing] stage SELECT (fetch baseline JSON): %.3fs", time.monotonic() - t1)
        row = cursor.fetchone()
        if row is None:
            LOGGER.warning("No baseline on stage for hash=%s", params_hash)
            return [], {}

        baseline: dict = json.loads(row[0])
        result_sets = baseline.get("result_sets", [[]])
        col_types_list: list[dict] = baseline.get("column_types", [{}])

        # Use the last result set for comparison, matching the Snowflake-native
        # path.  For simple procedures this is index 0 (the only set).  For
        # legacy capture baselines, index 0 is the OUT row and the last is the
        # data.  For template-driven baselines the last set is the final fetch.
        data_idx = len(result_sets) - 1 if result_sets else 0
        col_types: dict[str, str] = (
            col_types_list[data_idx] if data_idx < len(col_types_list) else {}
        )
        data_set: list = result_sets[data_idx] if data_idx < len(result_sets) else []
        if not col_types:
            return [], col_types
        if data_set and isinstance(data_set[0], dict):
            rows = [{k.upper(): v for k, v in r.items()} for r in data_set]
        else:
            rows = [dict(zip(col_types.keys(), r)) for r in data_set]

        # Coerce numeric columns to proper Python int/float so DuckDB compares
        # values numerically rather than lexicographically.
        numeric_cols = {col for col, typ in col_types.items() if is_numeric_type(typ)}
        rows = [
            {k: (coerce_numeric(v) if k in numeric_cols else v) for k, v in row.items()}
            for row in rows
        ]
        # Parse ISO timestamp strings to datetime objects so Polars infers the
        # same Datetime dtype as the Snowflake connector returns.
        timestamp_cols = {col for col, typ in col_types.items() if is_timestamp_type(typ)}
        rows = [
            {k: (coerce_timestamp(v) if k in timestamp_cols else v) for k, v in row.items()}
            for row in rows
        ]
        return rows, col_types
    finally:
        cursor.close()


# ---------------------------------------------------------------------------
# Fetch baseline output params into Python memory (for in-memory validation)
# ---------------------------------------------------------------------------

def fetch_baseline_output_rows(
    connector: ConnectorBase,
    database: str,
    params_hash: str,
) -> tuple[list[dict], dict[str, str]]:
    """Return ``(rows, col_types)`` for baseline output params.

    Reads staged baseline JSON and converts ``output_params`` to a one-row shape
    suitable for DuckDB in-memory comparison.
    """
    fq = f"{database}.{_SCHEMA}"
    stage = f"@{fq}.BASELINES"
    raw_fmt = f"{fq}.RAW_TEXT_FORMAT"

    cursor = connector.connection.cursor()
    try:
        cursor.execute(
            f"CREATE FILE FORMAT IF NOT EXISTS {raw_fmt} "
            f"TYPE = CSV FIELD_DELIMITER = NONE RECORD_DELIMITER = NONE"
        )
        cursor.execute(f"""\
SELECT $1 AS raw_text
FROM {stage} (FILE_FORMAT => '{raw_fmt}')
WHERE PARSE_JSON($1):params_hash::VARCHAR = '{params_hash}'
LIMIT 1""")
        row = cursor.fetchone()
        if row is None:
            LOGGER.warning("No baseline on stage for output params hash=%s", params_hash)
            return [], {}

        baseline: dict = json.loads(row[0])
        output_params: dict | None = baseline.get("output_params")
        if not output_params:
            return [], {}

        col_types: dict[str, str] = {
            k.upper(): v
            for k, v in (baseline.get("output_param_types") or {}).items()
        }
        if not col_types:
            col_types = {k.upper(): "VARCHAR" for k in output_params.keys()}

        numeric_cols = {col for col, typ in col_types.items() if is_numeric_type(typ)}
        row_data = {
            key.upper(): (coerce_numeric(value) if key.upper() in numeric_cols else value)
            for key, value in output_params.items()
        }
        return [row_data], col_types
    finally:
        cursor.close()


# ---------------------------------------------------------------------------
# Materialize output parameters -- baseline (from staged JSON)
# ---------------------------------------------------------------------------

def materialize_baseline_output_params_table(
    connector: ConnectorBase,
    database: str,
    procedure: str,
    params_hash: str,
) -> str | None:
    """Create a single-row transient table from baseline ``output_params``.

    Reads the raw JSON from the stage, extracts the ``output_params`` dict
    and the ``output_param_types`` dict, and creates a transient table with
    one column per output parameter cast to its original captured type.

    Returns ``None`` if the baseline has no output params.
    """
    table_name = _make_table_name(database, procedure, params_hash, "SOURCE_OUTPUTS")
    fq = f"{database}.{_SCHEMA}"
    stage = f"@{fq}.BASELINES"
    raw_fmt = f"{fq}.RAW_TEXT_FORMAT"

    cursor = connector.connection.cursor()
    try:
        cursor.execute(
            f"CREATE FILE FORMAT IF NOT EXISTS {raw_fmt} "
            f"TYPE = CSV FIELD_DELIMITER = NONE RECORD_DELIMITER = NONE"
        )
        cursor.execute(f"""\
SELECT $1 AS raw_text
FROM {stage} (FILE_FORMAT => '{raw_fmt}')
WHERE PARSE_JSON($1):params_hash::VARCHAR = '{params_hash}'
LIMIT 1""")
        row = cursor.fetchone()
        if row is None:
            return None

        baseline: dict = json.loads(row[0])
        output_params: dict | None = baseline.get("output_params")
        if not output_params:
            return None

        # Normalize to uppercase so the col_upper lookup below finds the correct type.
        stored_types: dict[str, str] = {
            k.upper(): v
            for k, v in (baseline.get("output_param_types") or {}).items()
        }

        col_expressions: list[str] = []
        for k, v in output_params.items():
            col_upper = k.upper()
            cast_type = stored_types.get(col_upper, "VARCHAR")
            if v is None:
                col_expressions.append(f"NULL::{cast_type} AS \"{col_upper}\"")
            elif isinstance(v, str):
                escaped = v.replace("'", "''")
                col_expressions.append(
                    f"'{escaped}'::{cast_type} AS \"{col_upper}\""
                )
            else:
                col_expressions.append(
                    f"{v}::{cast_type} AS \"{col_upper}\""
                )

        columns = ", ".join(col_expressions)
        cursor.execute(
            f"CREATE OR REPLACE TRANSIENT TABLE {table_name} AS SELECT {columns}"
        )
        return table_name
    finally:
        cursor.close()


# ---------------------------------------------------------------------------
# Materialize output parameters -- actual (via anonymous block)
# ---------------------------------------------------------------------------

def materialize_actual_output_params_table(
    connector: ConnectorBase,
    database: str,
    outputs_sql: str,
    procedure: str,
    params_hash: str,
) -> str:
    """Execute a Snowflake anonymous block and capture output params.

    The *outputs_sql* should be a rendered anonymous block (from the
    ``validate_outputs.sql.j2`` template) that returns output parameter
    values as a single-row result set.

    Returns the fully qualified transient table name.
    """
    table_name = _make_table_name(database, procedure, params_hash, "TARGET_OUTPUTS")

    cursor = connector.connection.cursor()
    try:
        cursor.execute(outputs_sql)
        query_id = cursor.sfqid
        cursor.execute(
            f"CREATE OR REPLACE TRANSIENT TABLE {table_name} "
            f"AS SELECT * FROM TABLE(RESULT_SCAN('{query_id}'))"
        )
    finally:
        cursor.close()

    return table_name


# ---------------------------------------------------------------------------
# Cleanup
# ---------------------------------------------------------------------------

def drop_transient_tables(connector: ConnectorBase, *table_names: str) -> None:
    """Drop transient tables, ignoring errors."""
    for name in table_names:
        try:
            connector.execute_statement(f"DROP TABLE IF EXISTS {name}")
        except Exception:
            pass
