"""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
from decimal import Decimal
from typing import TYPE_CHECKING
from pandas import DataFrame

if TYPE_CHECKING:
    import snowflake.snowpark
    from relationalai.semantics.backends.sql import sql_ir as sql
    from relationalai.semantics.backends.sql.sql_emitter import SQLEmitter

# 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',
}

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 execute_sql(self, sql: str) -> DataFrame:
        snow_df = self.session.sql(sql)
        df = snow_df.to_pandas()
        self._restore_numeric_columns(df, snow_df)
        return df

    def _restore_numeric_columns(self, pandas_df: DataFrame, snow_df: "snowflake.snowpark.DataFrame") -> None:
        """Restore numeric Snowflake columns that Snowpark materializes as float."""
        from snowflake.snowpark.types import (
            DecimalType,
            ByteType,
            ShortType,
            IntegerType,
            LongType,
        )

        # Snowpark schema field names can be quoted (e.g. '"Amount"') even when pandas columns are unquoted.
        by_name = {str(col): col for col in pandas_df.columns}
        for field in snow_df.schema.fields:
            raw_name = field.name
            unquoted_name = raw_name[1:-1] if raw_name.startswith('"') and raw_name.endswith('"') else raw_name
            col_name = by_name.get(raw_name) or by_name.get(unquoted_name)
            if col_name is None:
                continue

            if isinstance(field.datatype, (ByteType, ShortType, IntegerType, LongType)):
                series = pandas_df[col_name]
                if bool(series.isna().any()):
                    pandas_df[col_name] = series.astype("Int64")
                else:
                    pandas_df[col_name] = series.astype("int64")
                continue

            if not isinstance(field.datatype, DecimalType):
                continue

            scale = field.datatype.scale
            quant = Decimal(1).scaleb(-scale)
            def to_decimal(value):
                if value is None:
                    return None
                # pandas NaN / numpy NaN
                try:
                    if value != value:
                        return None
                except Exception:
                    pass
                if isinstance(value, Decimal):
                    return value.quantize(quant)
                return Decimal(f"{float(value):.{scale}f}").quantize(quant)

            pandas_df[col_name] = pandas_df[col_name].map(to_decimal)

    def execute_sql_batch(self, statements: list[str]) -> None:
        if not statements:
            return
        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.SemiNaiveLoop", emitter: "SQLEmitter") -> 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])

        def cols_sql(table: "sql.SemiNaiveTable") -> str:
            return ", ".join(emitter._quote_identifier(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"
            )

        cte_parts = []
        for table in loop.tables:
            cols = cols_by_table[table.table_name]
            store = store_names[table.table_name]
            cte_parts.append(
                f"{table.delta_view_name} AS (\n"
                f"  SELECT {cols} FROM {store} WHERE _gen = :iter\n"
                f")"
            )
            cte_parts.append(
                f"{table.accum_view_name} AS (\n"
                f"  SELECT {cols} FROM {store} WHERE _gen <= :iter\n"
                f")"
            )
        cte_prefix = "WITH " + ",\n".join(cte_parts)

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

            loop_body.append(
                f"delta_count := (SELECT COUNT(*) FROM ({next_delta_sql}) AS delta_cnt);\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_sql}) AS delta_rows;\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 = ".".join(emitter._quote_identifier(part) for part in formatted.split("."))
            final_stmts.append(
                f"CREATE OR REPLACE TEMP VIEW {table.accum_view_name} AS\n"
                f"SELECT {cols} FROM {store_name}"
            )
            final_stmts.append(
                f"CREATE OR REPLACE TABLE {view_name} AS\n{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()
