"""SQL-install-based executor: compiles the model to SQL views and installs them."""

from __future__ import annotations

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

from ...config.config import Config
from ...config import shims
from ...config.connections.duckdb import DuckDBConnection
from ...config.connections.snowflake import SnowflakeConnection
from .sql import SQLCompiler
from .sql.dialect import DuckDBDialect, SnowflakeDialect
from ...util.error import exc
from ...shims.v0debugging import span
from .executor import Executor
import rich

if TYPE_CHECKING:
    from ..frontend.base import Model, Fragment
    from .sql.emitter import SQL2Emitter
    from .sql import 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.SQL2SemiNaiveLoop", emitter: "SQL2Emitter") -> None:
        """Execute a semi-naive loop."""
        ...

#------------------------------------------------------
# SQLExecutor
#------------------------------------------------------

class SQLExecutor(Executor):
    def __init__(self, config: Config):
        self.config = config
        self._show_debug_logs = bool(getattr(config.debug, "show_debug_logs", False))
        self._installed_model_token: tuple[int, int] | None = None
        self._warned_model_token: tuple[int, 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.yaml 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 relationalai.semantics.backends.sql.duckdb_executor import DuckDBExecutor
            self.dialect_executor = DuckDBExecutor(conn.get_session())
            self.compiler = SQLCompiler(dialect=DuckDBDialect(), config=config)
        elif isinstance(conn, SnowflakeConnection):
            from relationalai.semantics.backends.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 _model_token(self, model: "Model") -> tuple[int, int]:
        return (id(model), model._version)

    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_query(self, model: "Model", fragment: "Fragment", **kwargs) -> DataFrame:
        return self.execute_dsl_query(model, fragment)

    def execute_dsl_query(self, model: "Model", fragment: "Fragment") -> DataFrame:
        self.compiler.compile_model(model)
        plan = self.compiler.execution_plan
        emitter = self.compiler.emitter
        dry_run = bool(shims.DRY_RUN or getattr(getattr(self.config, "compiler", None), "dry_run", False))
        model_token = self._model_token(model)
        install_needed = self._installed_model_token != model_token and plan.steps
        if install_needed and not dry_run:
            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 as e:
                        exc(
                            "Model install failed",
                            "Failed to execute model-install SQL batch.",
                            [str(e)],
                        )
                for loop in step.semi_naive_loops:
                    try:
                        self.dialect_executor.execute_semi_naive(loop, emitter)
                    except Exception as e:
                        exc(
                            "Model install failed",
                            "Failed to execute recursive model-install loop.",
                            [str(e)],
                        )
            self._installed_model_token = model_token
            self._emit_model_warnings_if_any(model)
        elif install_needed and dry_run:
            # Compile-only mode: generate install SQL for validation/logging without touching the backend.
            for step in plan.steps:
                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))
            self._installed_model_token = model_token

        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))
        if dry_run:
            return DataFrame()
        with span("query") as query_span:
            result = self.execute_sql(frag_sql)
            query_span["results"] = result
            query_span["sort_results"] = True
        return result

    def _emit_model_warnings_if_any(self, model: "Model") -> None:
        # Bridge Error entities to ModelWarning for SQL2 parity with non-SQL pipelines.
        pipeline = getattr(self.compiler, "pipeline", "")
        if pipeline != "sql2":
            return
        model_token = self._model_token(model)
        if self._warned_model_token == model_token:
            return

        analysis = getattr(self.compiler, "model_artifacts", None)
        model_analysis = getattr(analysis, "model_analysis", None)
        mm_model = getattr(model_analysis, "model", None)
        if mm_model is None:
            return

        error_type = next((t for t in mm_model.types if getattr(t, "name", None) == "Error"), None)
        if error_type is None:
            self._warned_model_token = model_token
            return

        message_rel = None
        for rel in mm_model.relations:
            if getattr(rel, "name", None) != "message":
                continue
            if len(getattr(rel, "fields", ())) < 2:
                continue
            key_type = rel.fields[0].type
            if getattr(key_type, "id", None) == getattr(error_type, "id", None):
                message_rel = rel
                break
        if message_rel is None:
            self._warned_model_token = model_token
            return

        value_ix = next((ix for ix, field in enumerate(message_rel.fields) if not field.input and ix != 0), None)
        if value_ix is None:
            value_ix = len(message_rel.fields) - 1

        warning_df = None
        analysis = getattr(self.compiler, "model_artifacts", None)
        emitter = getattr(self.compiler, "emitter", None)
        table_shape = getattr(analysis, "table_shape", None)
        resolve_schema = getattr(self.compiler, "_resolve_model_schema", None)
        if analysis is None or emitter is None or table_shape is None:
            return
        if not callable(resolve_schema):
            return
        schema_name = resolve_schema(model)
        table_name = table_shape.table_name_by_relation_id.get(message_rel.id)
        if not table_name:
            model_analysis = getattr(analysis, "model_analysis", None)
            if model_analysis is not None:
                table_info = model_analysis.get_table_info(message_rel)
                table_name = table_info.name if table_info is not None else None
        if not table_name:
            self._warned_model_token = model_token
            return

        value_col_name = table_shape.property_column_by_relation_id.get(message_rel.id)
        if not value_col_name:
            value_col_name = message_rel.fields[value_ix].name

        qid = getattr(emitter, "_qid", None)
        qid_path = getattr(emitter, "_qid_path", None)
        value_col = qid(value_col_name) if callable(qid) else f'"{value_col_name}"'
        table_fqn = table_name if schema_name is None else f"{schema_name}.{table_name}"
        rel_fqn = qid_path(table_fqn) if callable(qid_path) else table_fqn
        try:
            warning_df = self.execute_sql(f'SELECT DISTINCT {value_col} AS "_message" FROM {rel_fqn}')
        except Exception:
            return
        if warning_df is None:
            return

        if "_message" not in warning_df.columns:
            self._warned_model_token = model_token
            return

        try:
            from v0.relationalai.errors import ModelWarning
        except Exception:
            self._warned_model_token = model_token
            return

        for ix, raw in enumerate(warning_df["_message"].tolist()):
            if raw is None:
                continue
            message = str(raw)
            payload = [{
                "message": message,
                "props": {
                    "pyrel_id": f"{pipeline}_error_{model._version}_{ix}",
                    "message": message,
                },
            }]
            warnings.warn(ModelWarning(payload), stacklevel=3)
        self._warned_model_token = model_token
