"""DuckDB-specific executor implementation.

This module provides the DuckDBExecutor class that implements the Executor
protocol for DuckDB. It handles:
- Schema fetching via INFORMATION_SCHEMA
- DuckDB-specific table name formatting (db.schema.table → db_schema.table)
- Semi-naive iteration helpers for recursive plans
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any
import pandas as pd
from pandas import DataFrame
import pyarrow as pa
from v0.relationalai.clients.result_helpers import Int128Dtype

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

# Type mapping from DuckDB types to Python types
DUCKDB_TYPE_MAP: dict[str, type | str] = {
    'VARCHAR': str,
    'TEXT': str,
    'CHAR': str,
    'STRING': str,
    'INTEGER': int,
    'BIGINT': int,
    'SMALLINT': int,
    'TINYINT': int,
    'INT': int,
    'INT4': int,
    'INT8': int,
    'HUGEINT': int,
    'DOUBLE': float,
    'FLOAT': float,
    'REAL': float,
    'FLOAT4': float,
    'FLOAT8': float,
    'BOOLEAN': bool,
    'BOOL': bool,
    'DATE': 'Date',
    'TIMESTAMP': 'DateTime',
    'TIMESTAMPTZ': 'DateTime',
    'TIMESTAMP WITH TIME ZONE': 'DateTime',
    'TIME': 'DateTime',
    'DECIMAL': 'Decimal',
    'NUMERIC': 'Decimal',
}

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

_SIGNED_INT_DTYPES = {
    False: {8: "int8", 16: "int16", 32: "int32", 64: "int64"},
    True: {8: "Int8", 16: "Int16", 32: "Int32", 64: "Int64"},
}
_UNSIGNED_INT_DTYPES = {
    False: {8: "uint8", 16: "uint16", 32: "uint32", 64: "uint64"},
    True: {8: "UInt8", 16: "UInt16", 32: "UInt32", 64: "UInt64"},
}

def format_duckdb_columns(result_frame: pd.DataFrame, arrow_schema: pa.Schema) -> pd.DataFrame:
    for field in arrow_schema:
        col = field.name
        series = result_frame[col]
        mask = series.isna()
        if bool(mask.all()):
            result_frame[col] = series.astype("float64")
            continue
        any_null = bool(mask.any())
        pa_type = field.type
        if pa.types.is_decimal(pa_type):
            if pa_type.scale == 0:  # integer-like
                series = _cast_integer_column(series, pa_type.precision, any_null=any_null)
        elif pa.types.is_integer(pa_type):
            series = _cast_arrow_integer_column(series, pa_type, any_null=any_null)
        result_frame[col] = series
    return result_frame

def _cast_integer_column(series: pd.Series, precision: int, *, any_null: bool) -> pd.Series:
    """Cast a numeric pandas Series to Int64 or Int128 depending on precision."""
    if any_null:
        return series.astype(Int128Dtype())
    if precision <= 19:
        return series.astype("int64")
    return series.astype(Int128Dtype())

def _cast_arrow_integer_column(series: pd.Series, pa_type: pa.DataType, *, any_null: bool) -> pd.Series:
    """Cast a pandas Series to a nullable integer dtype that matches the Arrow type."""
    if any_null:
        return series.astype(Int128Dtype())
    bit_width = getattr(pa_type, "bit_width", 64)
    if bit_width > 64:
        return series.astype(Int128Dtype())
    if pa.types.is_unsigned_integer(pa_type):
        unsigned_dtype: Any = _UNSIGNED_INT_DTYPES[False].get(bit_width, "uint64")
        return series.astype(unsigned_dtype)
    signed_dtype: Any = _SIGNED_INT_DTYPES[False].get(bit_width, "int64")
    return series.astype(signed_dtype)

#------------------------------------------------------
# Executor
#------------------------------------------------------

class DuckDBExecutor:
    dialect = "duckdb"

    def __init__(
        self,
        conn: Any,
    ):
        self.conn = conn

    @staticmethod
    def format_table_fqn(name: str, is_temp: bool = False) -> str:
        parts = name.split(".")
        if is_temp:
            # Temp tables in DuckDB don't support schema prefixes
            return parts[-1]
        if len(parts) == 3:
            return f"{parts[0]}_{parts[1]}.{parts[2]}"
        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]:
        # DuckDB maps database.schema → database_schema as the schema name
        duckdb_schema = f"{database}_{schema}"
        query = f"""
            SELECT column_name, data_type
            FROM information_schema.columns
            WHERE table_schema = '{duckdb_schema}'
              AND table_name = '{table}'
            ORDER BY ordinal_position
        """
        try:
            result = self.conn.execute(query).fetchall()
            return {row[0]: self._map_duckdb_type(row[1]) for row in result}
        except Exception:
            # If INFORMATION_SCHEMA query fails, return empty dict
            return {}

    def _map_duckdb_type(self, dtype: str) -> type | str:
        # Extract base type (remove precision/scale info)
        base_type = dtype.upper().split('(')[0].strip()
        return DUCKDB_TYPE_MAP.get(base_type, str)

    def execute_sql(self, sql: str) -> DataFrame:
        """Execute a SQL query and return results.

        Args:
            sql: The SQL query string to execute

        Returns:
            A pandas DataFrame containing the query results
        """
        arrow_table = self.conn.execute(sql).fetch_arrow_table()

        def _types_mapper(pa_type: pa.DataType):
            if pa.types.is_decimal(pa_type):
                return pd.ArrowDtype(pa_type)
            return None

        try:
            df = arrow_table.to_pandas(types_mapper=_types_mapper)
            return format_duckdb_columns(df, arrow_table.schema)
        except ValueError as exc:
            # PyArrow rejects duplicate column labels in some queries.
            # Build the frame manually so we preserve duplicate aliases.
            if "non-unique column index" not in str(exc):
                raise
            data_by_position = {
                f"c_{ix}": arrow_table.column(ix).to_pylist()
                for ix in range(arrow_table.num_columns)
            }
            df = pd.DataFrame(data_by_position)
            df.columns = list(arrow_table.column_names)
            return df

    def execute_sql_batch(self, statements: list[str]) -> None:
        if not statements:
            return
        script = ";\n".join(stmt.rstrip().rstrip(";") for stmt in statements) + ";"
        self.conn.execute(script)

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

    def execute_semi_naive(self, loop: "sql.SQL2SemiNaiveLoop", emitter: "SQL2Emitter") -> None:
        max_iterations = 1000
        conn = self.conn

        def _qid(name: str) -> str:
            return emitter._quote_identifier(name)

        def _temp_table(name: str, query: str) -> None:
            conn.execute(f"CREATE OR REPLACE TEMP TABLE {_qid(name)} AS\n{query}")

        def _swap_table(name: str, next_name: str) -> None:
            conn.execute(f"DROP TABLE IF EXISTS {_qid(name)}")
            conn.execute(f"ALTER TABLE {_qid(next_name)} RENAME TO {_qid(name)}")

        def _append_table_suffix(path: str, suffix: str) -> str:
            if "." not in path:
                return f"{path}{suffix}"
            prefix, table = path.rsplit(".", 1)
            return f"{prefix}.{table}{suffix}"

        def _state_table_path(public_path: str) -> str:
            # Keep recursive state internal and deterministic without reusing user-facing names.
            return _append_table_suffix(public_path, "__sql2_rec_state")

        # Initialize accum and delta tables
        for table in loop.tables:
            base_sql = emitter.emit(table.base_query)
            _temp_table(table.accum_view_name, base_sql)
            _temp_table(table.delta_view_name, base_sql)

        for _ in range(max_iterations):
            total_rows = 0
            # Compute next deltas for all tables using current deltas
            for table in loop.tables:
                delta_sql = emitter.emit(table.delta_query)
                _temp_table(
                    table.next_delta_view_name,
                    f"{delta_sql}\nEXCEPT\nSELECT * FROM {_qid(table.accum_view_name)}",
                )
                count_row = conn.execute(f"SELECT COUNT(*) FROM {_qid(table.next_delta_view_name)}").fetchone()
                if count_row is not None:
                    total_rows += int(count_row[0])

            if total_rows == 0:
                break

            # Merge and advance deltas
            for table in loop.tables:
                merge_sql = emitter.emit(table.merge_query)
                accum_next = f"{table.accum_view_name}__next"
                _temp_table(accum_next, merge_sql)
                _swap_table(table.accum_view_name, accum_next)
                _swap_table(table.delta_view_name, table.next_delta_view_name)

        # Finalize views
        for table in loop.tables:
            final_sql = emitter.emit(table.final_query)
            formatted = emitter.dialect.format_table_name(table.table_name)
            view_name = emitter._qid_path(formatted)
            state_table_name = _state_table_path(formatted)
            state_table_fqn = emitter._qid_path(state_table_name)
            conn.execute(f"CREATE OR REPLACE TABLE {state_table_fqn} AS\n{final_sql}")
            conn.execute(f"CREATE OR REPLACE VIEW {view_name} AS\nSELECT * FROM {state_table_fqn}")
            conn.execute(f"DROP TABLE IF EXISTS {_qid(table.accum_view_name)}")
            conn.execute(f"DROP TABLE IF EXISTS {_qid(table.delta_view_name)}")
            conn.execute(f"DROP TABLE IF EXISTS {_qid(table.next_delta_view_name)}")
