"""Executor Protocol for database-agnostic execution.

This module defines the Executor protocol that both DuckDBExecutor and
SnowflakeExecutor implement, providing a unified interface for:
- Running SQL queries
- Fetching schema information
- Formatting table names for the target database
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Protocol, runtime_checkable
from pandas import DataFrame
from sqlglot import parse_one

from relationalai.config.config import Config, create_config
from relationalai.config.connections.duckdb import DuckDBConnection
from relationalai.config.connections.snowflake import SnowflakeConnection
from relationalai.semantics.backends.sql.sql_compiler import SQLCompiler
from relationalai.semantics.backends.sql.sql_emitter import DuckDBDialect, SnowflakeDialect
from relationalai.util.error import exc
from ...shims.v0debugging import span
import rich

if TYPE_CHECKING:
    from ..frontend.base import Model, Fragment
    from relationalai.semantics.backends.sql.sql_emitter import SQLEmitter
    from relationalai.semantics.backends.sql import sql_ir as sql

#------------------------------------------------------
# DialectExecutor
#------------------------------------------------------

@runtime_checkable
class DialectExecutor(Protocol):
    dialect: str

    def fetch_schema(self, database: str, schema: str, table: str) -> dict[str, type | str]:
        """Fetch column types for a table.

        Args:
            database: The database name
            schema: The schema name
            table: The table name

        Returns:
            A dictionary mapping column names to Python types
        """
        ...

    def table_fqn(self, name: str) -> str:
        """Format a table name for this database.

        Different databases have different naming conventions:
        - DuckDB: db.schema.table → db_schema.table (underscore naming)
        - Snowflake: db.schema.table → db.schema.table (dot notation, uppercased)

        Args:
            name: The fully qualified table name

        Returns:
            The formatted table name for this database
        """
        ...

    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
        """
        ...

    def execute_sql_batch(self, statements: list[str]) -> None:
        """Execute multiple SQL statements in a batch."""
        ...

    def execute_semi_naive(self, loop: "sql.SemiNaiveLoop", emitter: "SQLEmitter") -> None:
        """Execute a semi-naive loop."""
        ...

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

class Executor:
    def __init__(self, config: Config = create_config()):
        self.config = config
        self._show_debug_logs = bool(getattr(config.debug, "show_debug_logs", False))
        self._installed_model_version: int | None = None
        self._schema_ensured: bool = False
        self._configured_schema: str
        self.dialect_executor: DialectExecutor
        self.compiler: SQLCompiler
        conn = config.get_default_connection()
        schema = config.model.schema_
        if not isinstance(schema, str) or not schema.strip():
            exc("Missing config", "SQL execution requires `model.schema` to be configured in order to install the model.", [
                "Set it in your raiconfig.toml to tell RelationalAI where to install the views and tables for the model.",
                "[cyan]model:\n"
                "[cyan]  schema: your.schema",
            ])
        self._configured_schema = schema.strip()
        if isinstance(conn, DuckDBConnection):
            from .sql.duckdb_executor import DuckDBExecutor
            self.dialect_executor = DuckDBExecutor(conn.get_session())
            self.compiler = SQLCompiler(dialect=DuckDBDialect(), config=config)
        elif isinstance(conn, SnowflakeConnection):
            from .sql.snowflake_executor import SnowflakeExecutor
            self.dialect_executor = SnowflakeExecutor(conn.get_session())
            self.compiler = SQLCompiler(dialect=SnowflakeDialect(), config=config)
        else:
            raise ValueError(f"Unsupported connection type: {type(config.default_connection)}")

    def _schema_name_for_create(self) -> str:
        schema = self._configured_schema
        if self.dialect_executor.dialect == "duckdb":
            parts = schema.split(".")
            if len(parts) == 2:
                return f"{parts[0]}_{parts[1]}"
        return schema

    def _ensure_schema_exists(self) -> None:
        if self._schema_ensured:
            return
        schema_name = self._schema_name_for_create()
        try:
            self.execute_sql(f"CREATE SCHEMA IF NOT EXISTS {schema_name}")
        except Exception as e:
            exc(
                "Schema setup failed",
                f"Failed to ensure install schema '{self._configured_schema}' exists.",
                [str(e)],
            )
        self._schema_ensured = True

    def execute_sql(self, sql: str) -> DataFrame:
        return self.dialect_executor.execute_sql(sql)

    def _pretty_sql(self, sql: str) -> str:
        if not self._show_debug_logs:
            return sql
        dialect = "snowflake" if self.dialect_executor.dialect == "snowflake" else "duckdb"
        try:
            return parse_one(sql, read=dialect).sql(dialect=dialect, pretty=True)
        except Exception:
            return sql

    def execute_dsl_query(self, model: Model, fragment: Fragment) -> DataFrame:
        self.compiler.compile_model(model)
        plan = self.compiler.execution_plan
        emitter = self.compiler.emitter
        install_needed = self._installed_model_version != model._version and plan.steps
        install_ok = True
        if install_needed:
            self._ensure_schema_exists()
            for step in plan.steps:
                view_sql_batch: list[str] = []
                for view in step.views:
                    sql_str = emitter.emit(view)
                    if self._show_debug_logs:
                        rich.print("\n[yellow]--------------------------------------------------------------[/yellow]")
                        rich.print(self._pretty_sql(sql_str))
                    view_sql_batch.append(sql_str)
                if view_sql_batch:
                    try:
                        self.dialect_executor.execute_sql_batch(view_sql_batch)
                    except Exception:
                        # Fallback to one-by-one so we can surface which statement failed.
                        for sql_str in view_sql_batch:
                            try:
                                self.execute_sql(sql_str)
                            except Exception as e:
                                install_ok = False
                                print(f"  Warning: Failed to execute model statement: {e}")
                for loop in step.semi_naive_loops:
                    try:
                        self.dialect_executor.execute_semi_naive(loop, emitter)
                    except Exception as e:
                        install_ok = False
                        print(f"  Warning: Failed to execute recursive loop: {e}")
            if install_ok:
                self._installed_model_version = model._version

        frag_sql = self.compiler.compile_fragment(fragment)
        if frag_sql is None:
            return DataFrame()
        if self._show_debug_logs:
            rich.print("\n[cyan]--------------------------------------------------------------[/cyan]")
            rich.print(self._pretty_sql(frag_sql))
        with span("query") as query_span:
            result = self.execute_sql(frag_sql)
            query_span["results"] = result
            query_span["sort_results"] = True
        return result
