"""Snowflake-specific executor implementation.

This module provides the SnowflakeExecutor class that implements the Executor
protocol for Snowflake. It handles:
- Schema fetching via SHOW COLUMNS
- Snowflake-specific table name formatting (uppercase, dot notation)
- Semi-naive iteration using Snowflake Scripting
"""

from __future__ import annotations

import json
import re
from decimal import Decimal
from math import isnan
from typing import TYPE_CHECKING, Any, cast
import pandas as pd
from pandas import DataFrame
from v0.relationalai.clients.result_helpers import Int128Dtype

if TYPE_CHECKING:
    import snowflake.snowpark
    from relationalai.semantics.backends.sql import ir as sql
    from relationalai.semantics.backends.sql.emitter import SQL2Emitter

# Type mapping from Snowflake types to Python/RAI types
SNOWFLAKE_TYPE_MAP: dict[str, type | str] = {
    'TEXT': str,
    'VARCHAR': str,
    'CHAR': str,
    'STRING': str,
    'BINARY': bytes,
    'NUMBER': 'Decimal',
    'FIXED': 'Decimal',
    'FLOAT': float,
    'REAL': float,
    'DOUBLE': float,
    'BOOLEAN': bool,
    'DATE': 'Date',
    'TIME': 'DateTime',
    'TIMESTAMP': 'DateTime',
    'TIMESTAMP_LTZ': 'DateTime',
    'TIMESTAMP_NTZ': 'DateTime',
    'TIMESTAMP_TZ': 'DateTime',
}


def _normalize_column_name(name: str) -> str:
    if name.startswith('"') and name.endswith('"'):
        return name[1:-1]
    return name


def _collect_sql_output_types(
    snow_df: "snowflake.snowpark.DataFrame",
    column_names: list[str],
) -> tuple[dict[str, int | str | None], ...]:
    out: list[dict[str, int | str | None]] = []
    for ix, field in enumerate(snow_df.schema.fields):
        col_name = column_names[ix] if ix < len(column_names) else _normalize_column_name(field.name)
        out.append(
            {
                "name": col_name,
                "snowflake_type": type(field.datatype).__name__,
                "precision": getattr(field.datatype, "precision", None),
                "scale": getattr(field.datatype, "scale", None),
            }
        )
    return tuple(out)


def format_snowflake_columns(
    result_frame: pd.DataFrame,
    snow_df: "snowflake.snowpark.DataFrame",
) -> pd.DataFrame:
    """Schema-driven DataFrame formatting for Snowflake results."""
    if result_frame.shape[0] == 0:
        return result_frame

    from snowflake.snowpark.types import (
        DecimalType,
        ByteType,
        ShortType,
        IntegerType,
        LongType,
        StringType,
    )

    int64_min = -(1 << 63)
    int64_max = (1 << 63) - 1

    def _is_null(value: object) -> bool:
        if value is None or value is pd.NA:
            return True
        if isinstance(value, float):
            try:
                return isnan(value)
            except Exception:
                return False
        return False

    def _requires_int128(series: pd.Series) -> bool:
        for value in series:
            if _is_null(value):
                continue
            int_value = int(value)
            if int_value < int64_min or int_value > int64_max:
                return True
        return False

    def _cast_integer_series(series: pd.Series) -> pd.Series:
        if _requires_int128(series):
            return series.astype(Int128Dtype())
        if bool(series.isna().any()):
            return series.astype("Int64")
        return series.astype("int64")

    def _to_decimal(value: object, *, scale: int) -> Decimal | None:
        if _is_null(value):
            return None
        quant = Decimal(1).scaleb(-scale)
        if isinstance(value, Decimal):
            return value.quantize(quant)
        return Decimal(str(value)).quantize(quant)

    def _cast_bool_string_series(series: pd.Series) -> pd.Series:
        normalized = series.map(lambda v: None if _is_null(v) else str(v).strip().lower())
        valid = {v for v in normalized.tolist() if v is not None}
        if not valid.issubset({"true", "false"}):
            return series
        return normalized.map(lambda v: None if v is None else (v == "true")).astype("boolean")

    for field_ix, field in enumerate(snow_df.schema.fields):
        if field_ix >= len(result_frame.columns):
            break
        series = result_frame.iloc[:, field_ix]
        def _assign(series_value: pd.Series) -> None:
            # Use positional assignment so extension dtypes (e.g. Int128) survive.
            result_frame.isetitem(field_ix, cast(Any, series_value))

        if isinstance(field.datatype, (ByteType, ShortType, IntegerType, LongType)):
            _assign(_cast_integer_series(series))
            continue

        if not isinstance(field.datatype, DecimalType):
            if isinstance(field.datatype, StringType):
                _assign(_cast_bool_string_series(series))
            continue

        scale = field.datatype.scale
        if scale == 0:
            _assign(_cast_integer_series(series))
            continue
        _assign(series.map(lambda value: _to_decimal(value, scale=scale)))

    return result_frame


class SnowflakeExecutor:
    dialect = "snowflake"

    def __init__(
        self,
        session: "snowflake.snowpark.Session",
    ):
        self.session = session

    @staticmethod
    def format_table_fqn(name: str, is_temp: bool = False) -> str:
        return name

    def table_fqn(self, name: str, is_temp: bool = False) -> str:
        return self.format_table_fqn(name, is_temp=is_temp)

    def fetch_schema(self, database: str, schema: str, table: str) -> dict[str, type | str]:
        try:
            # Use SHOW COLUMNS which works with Snowflake permissions
            show_query = f"SHOW COLUMNS IN {database}.{schema}.{table}"
            self.session.sql(show_query).collect()

            # Get results from result_scan
            result_query = """
                SELECT "column_name", "data_type"
                FROM TABLE(RESULT_SCAN(LAST_QUERY_ID()))
            """
            rows = self.session.sql(result_query).collect()

            result: dict[str, type | str] = {}
            for row in rows:
                column_name = str(row[0])
                data_type_json = row[1]
                if not isinstance(data_type_json, (str, bytes, bytearray)):
                    continue

                # Parse the JSON data type info
                type_info = json.loads(data_type_json)
                sf_type = type_info.get("type", "TEXT")

                result[column_name] = self._map_snowflake_type(sf_type, type_info)

            return result
        except Exception:
            # If schema fetch fails, return empty dict
            return {}

    def _map_snowflake_type(self, sf_type: str, type_info: dict) -> type | str:
        base_type = sf_type.upper()

        # Handle FIXED/NUMBER with precision/scale
        if base_type in ('FIXED', 'NUMBER'):
            precision = type_info.get('precision', 38)
            scale = type_info.get('scale', 0)

            if scale == 0:
                # Integer types
                if precision <= 18:
                    return 'Decimal(18,0)'
                else:
                    return 'Decimal(38,0)'
            else:
                return f'Decimal({precision},{scale})'

        return SNOWFLAKE_TYPE_MAP.get(base_type, str)

    def _create_table_target(self, sql: str) -> str | None:
        match = re.match(r'^\s*CREATE\s+OR\s+REPLACE\s+TABLE\s+(.+?)\s+AS\b', sql, flags=re.IGNORECASE | re.DOTALL)
        if not match:
            return None
        return match.group(1).strip()

    def _drop_view_if_exists_for_create_table(self, sql: str) -> None:
        target = self._create_table_target(sql)
        if not target:
            return
        try:
            self.session.sql(f"DROP VIEW IF EXISTS {target}").collect()
        except Exception:
            pass

    def execute_sql(self, sql: str) -> DataFrame:
        self._drop_view_if_exists_for_create_table(sql)
        snow_df = self.session.sql(sql)
        column_names = [_normalize_column_name(name) for name in snow_df.columns]
        rows = [tuple(row) for row in snow_df.collect()]
        df = DataFrame(rows, columns=column_names)
        formatted_df = format_snowflake_columns(df, snow_df)
        formatted_df.attrs["sql_output_types"] = _collect_sql_output_types(snow_df, column_names)
        return formatted_df

    def execute_sql_batch(self, statements: list[str]) -> None:
        if not statements:
            return
        for statement in statements:
            self._drop_view_if_exists_for_create_table(statement)
        body = "\n".join(stmt.rstrip().rstrip(";") + ";" for stmt in statements)
        script = "BEGIN\n" + body + "\nEND;"
        self.session.sql(script).collect()

    #------------------------------------------------------
    # Semi-naive execution
    #------------------------------------------------------

    def execute_semi_naive(self, loop: "sql.SQL2SemiNaiveLoop", emitter: "SQL2Emitter") -> None:
        max_iterations = 1000
        schema_context: str | None = None
        if loop.tables:
            first_table_name = loop.tables[0].table_name
            parts = first_table_name.split(".")
            if len(parts) >= 3:
                schema_context = ".".join(parts[:2])

        reserved_identifiers = set(getattr(emitter.dialect, "_RESERVED_IDENTIFIERS", ()))

        def qcol(name: str) -> str:
            token = name.strip()
            if (
                re.fullmatch(r"[A-Za-z_][A-Za-z0-9_$]*", token)
                and token.upper() == token
                and token.upper() not in reserved_identifiers
            ):
                return token
            return emitter._quote_identifier(token)

        def cols_sql(table: "sql.SQL2SemiNaiveTable") -> str:
            return ", ".join(qcol(c) for c in table.columns)

        # Build initialization statements
        init_stmts = []
        store_names = {table.table_name: f"{table.accum_view_name}__store" for table in loop.tables}
        cols_by_table = {table.table_name: cols_sql(table) for table in loop.tables}
        for table in loop.tables:
            store_name = store_names[table.table_name]
            base_sql = emitter.emit(table.base_query)
            cols = cols_by_table[table.table_name]
            init_stmts.append(
                f"CREATE OR REPLACE TEMP TABLE {store_name} AS\n"
                f"SELECT {cols}, 0 AS _gen\n"
                f"FROM ({base_sql}) AS base"
            )

        # Build loop body per table
        loop_body = []
        for table in loop.tables:
            store_name = store_names[table.table_name]
            cols = cols_by_table[table.table_name]
            delta_sql = emitter.emit(table.delta_query)
            next_delta_name = table.next_delta_view_name
            next_delta_sql = (
                f"SELECT {cols} FROM ({delta_sql}) AS delta_src\n"
                f"EXCEPT\n"
                f"SELECT {cols} FROM {store_name}"
            )

            loop_body.append(
                f"CREATE OR REPLACE TEMP TABLE {table.delta_view_name} AS\n"
                f"SELECT {cols} FROM {store_name} WHERE _gen = :iter;\n"
                f"CREATE OR REPLACE TEMP TABLE {table.accum_view_name} AS\n"
                f"SELECT {cols} FROM {store_name} WHERE _gen <= :iter;\n"
                f"CREATE OR REPLACE TEMP TABLE {next_delta_name} AS\n"
                f"{next_delta_sql};\n"
                f"delta_count := (SELECT COUNT(*) FROM {next_delta_name});\n"
                f"total_rows := total_rows + delta_count;\n"
                f"IF (delta_count > 0) THEN\n"
                f"  INSERT INTO {store_name} ({cols}, _gen)\n"
                f"  SELECT {cols}, :iter + 1 FROM {next_delta_name};\n"
                f"END IF;"
            )

        # Final views
        final_stmts = []
        for table in loop.tables:
            store_name = store_names[table.table_name]
            cols = cols_by_table[table.table_name]
            final_sql = emitter.emit(table.final_query)
            formatted = emitter.dialect.format_table_name(table.table_name)
            view_name = emitter._qid_path(formatted)
            final_stmts.append(
                f"CREATE OR REPLACE TEMP TABLE {table.accum_view_name} AS\n"
                f"SELECT {cols} FROM {store_name}"
            )
            final_stmts.append(
                f"CREATE OR REPLACE TABLE {view_name} AS\n"
                f"{final_sql}"
            )

        script = "\n".join([
            "DECLARE",
            "  iter INTEGER := 0;",
            "  total_rows INTEGER := 1;",
            "  delta_count INTEGER := 0;",
            "  max_iter INTEGER := " + str(max_iterations) + ";",
            "BEGIN",
            *([f"  USE SCHEMA {schema_context};"] if schema_context else []),
            *[f"  {stmt};" for stmt in init_stmts],
            "  WHILE (total_rows > 0 AND iter < max_iter) DO",
            "    total_rows := 0;",
            *["    " + line for stmt in loop_body for line in stmt.splitlines()],
            "    iter := iter + 1;",
            "  END WHILE;",
            *[f"  {stmt};" for stmt in final_stmts],
            "END;",
        ])

        self.session.sql(script).collect()
