"""Ibis-based metadata store for SQL databases.

Supports any SQL database that Ibis supports:
- DuckDB, PostgreSQL, MySQL (local/embedded)
- ClickHouse, Snowflake, BigQuery (cloud analytical)
- And 20+ other backends
"""

from abc import ABC, abstractmethod
from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any

import narwhals as nw
from narwhals.typing import Frame
from pydantic import Field

from metaxy._decorators import public
from metaxy.metadata_store.base import (
    MetadataStore,
    MetadataStoreConfig,
    VersioningEngineOptions,
)
from metaxy.metadata_store.exceptions import (
    FeatureNotFoundError,
    HashAlgorithmNotSupportedError,
    StoreNotOpenError,
    TableNotFoundError,
)
from metaxy.metadata_store.types import AccessMode
from metaxy.models.plan import FeaturePlan
from metaxy.models.types import CoercibleToFeatureKey, FeatureKey
from metaxy.versioning.ibis import IbisVersioningEngine
from metaxy.versioning.types import HashAlgorithm

if TYPE_CHECKING:
    import ibis
    import ibis.expr.types
    from ibis.backends.sql import SQLBackend


class IbisMetadataStoreConfig(MetadataStoreConfig):
    """Configuration for IbisMetadataStore.

    Example:
        ```toml title="metaxy.toml"
        [stores.dev.config]
        connection_string = "postgresql://user:pass@host:5432/db"
        table_prefix = "prod_"
        ```
    """

    connection_string: str | None = Field(
        default=None,
        description="Ibis connection string (e.g., 'clickhouse://host:9000/db').",
    )

    backend: str | None = Field(
        default=None,
        description="Ibis backend name (e.g., 'clickhouse', 'postgres', 'duckdb').",
        json_schema_extra={"mkdocs_metaxy_hide": True},
    )

    connection_params: dict[str, Any] | None = Field(
        default=None,
        description="Backend-specific connection parameters.",
    )

    table_prefix: str | None = Field(
        default=None,
        description="Optional prefix for all table names.",
    )

    auto_create_tables: bool | None = Field(
        default=None,
        description="If True, create tables on open. For development/testing only.",
    )


@public
class IbisMetadataStore(MetadataStore, ABC):
    """
    Generic SQL metadata store using Ibis.

    Supports any Ibis backend that supports struct types, such as: DuckDB, PostgreSQL, ClickHouse, and others.

    Warning:
        Backends without native struct support (e.g., SQLite) are NOT supported.

    Storage layout:
    - Each feature gets its own table: {feature}__{key}
    - System tables: metaxy__system__feature_versions, metaxy__system__events
    - Uses Ibis for cross-database compatibility

    Note: Uses MD5 hash by default for cross-database compatibility.
    DuckDBMetadataStore overrides this with dynamic algorithm detection.
    For other backends, override the calculator instance variable with backend-specific implementations.

    Example:
        <!-- skip next -->
        ```py
        # ClickHouse
        store = IbisMetadataStore("clickhouse://user:pass@host:9000/db")

        # PostgreSQL
        store = IbisMetadataStore("postgresql://user:pass@host:5432/db")

        # DuckDB (use DuckDBMetadataStore instead for better hash support)
        store = IbisMetadataStore("duckdb:///metadata.db")

        with store.open("w"):
            store.write(MyFeature, df)
        ```
    """

    versioning_engine_cls = IbisVersioningEngine

    def __init__(
        self,
        versioning_engine: VersioningEngineOptions = "auto",
        connection_string: str | None = None,
        *,
        backend: str | None = None,
        connection_params: dict[str, Any] | None = None,
        table_prefix: str | None = None,
        **kwargs: Any,
    ):
        """
        Initialize Ibis metadata store.

        Args:
            versioning_engine: Which versioning engine to use.
                - "auto": Prefer the store's native engine, fall back to Polars if needed
                - "native": Always use the store's native engine, raise `VersioningEngineMismatchError`
                    if provided dataframes are incompatible
                - "polars": Always use the Polars engine
            connection_string: Ibis connection string (e.g., "clickhouse://host:9000/db")
                If provided, backend and connection_params are ignored.
            backend: Ibis backend name (e.g., "clickhouse", "postgres", "duckdb")
                Used with connection_params for more control.
            connection_params: Backend-specific connection parameters
                e.g., {"host": "localhost", "port": 9000, "database": "default"}
            table_prefix: Optional prefix applied to all feature and system table names.
                Useful for logically separating environments (e.g., "prod_"). Must form a valid SQL
                identifier when combined with the generated table name.
            **kwargs: Passed to MetadataStore.__init__ (e.g., fallback_stores, hash_algorithm)

        Raises:
            ValueError: If neither connection_string nor backend is provided
            ImportError: If Ibis or required backend driver not installed

        Example:
            <!-- skip next -->
            ```py
            # Using connection string
            store = IbisMetadataStore("clickhouse://user:pass@host:9000/db")

            # Using backend + params
            store = IbisMetadataStore(backend="clickhouse", connection_params={"host": "localhost", "port": 9000})
            ```
        """
        from ibis.backends.sql import SQLBackend

        self.connection_string = connection_string
        self.backend = backend
        self.connection_params = connection_params or {}
        self._conn: SQLBackend | None = None
        self._table_prefix = table_prefix or ""

        super().__init__(
            **kwargs,
            versioning_engine=versioning_engine,
        )

    def _has_feature_impl(self, feature: CoercibleToFeatureKey) -> bool:
        feature_key = self._resolve_feature_key(feature)
        table_name = self.get_table_name(feature_key)
        return table_name in self.conn.list_tables()

    def get_table_name(
        self,
        key: FeatureKey,
    ) -> str:
        """Generate the storage table name for a feature or system table.

        Applies the configured table_prefix (if any) to the feature key's table name.
        Subclasses can override this method to implement custom naming logic.

        Args:
            key: Feature key to convert to storage table name.

        Returns:
            Storage table name with optional prefix applied.
        """
        base_name = key.table_name

        return f"{self._table_prefix}{base_name}" if self._table_prefix else base_name

    def _get_default_hash_algorithm(self) -> HashAlgorithm:
        """Get default hash algorithm for Ibis stores.

        Uses MD5 as it's universally supported across SQL databases.
        Subclasses like DuckDBMetadataStore can override for better algorithms.
        """
        return HashAlgorithm.MD5

    @contextmanager
    def _create_versioning_engine(self, plan: FeaturePlan) -> Iterator[IbisVersioningEngine]:
        """Create provenance engine for Ibis backend as a context manager.

        Args:
            plan: Feature plan for the feature we're tracking provenance for

        Yields:
            IbisVersioningEngine with backend-specific hash functions.

        Note:
            Base implementation only supports MD5 (universally available).
            Subclasses can override _create_hash_functions() for backend-specific hashes.
        """
        if self._conn is None:
            raise RuntimeError(
                "Cannot create provenance engine: store is not open. Ensure store is used as context manager."
            )

        # Create hash functions for Ibis expressions
        hash_functions = self._create_hash_functions()

        # Create engine using the configured class (allows subclass override)
        engine = self.versioning_engine_cls(
            plan=plan,
            hash_functions=hash_functions,  # ty: ignore[unknown-argument]
        )

        try:
            yield engine  # ty: ignore[invalid-yield]
        finally:
            # No cleanup needed for Ibis engine
            pass

    @abstractmethod
    def _create_hash_functions(self):
        """Create hash functions for Ibis expressions.

        Base implementation returns empty dict. Subclasses must override
        to provide backend-specific hash function implementations.

        Returns:
            Dictionary mapping HashAlgorithm to Ibis expression functions
        """
        return {}

    def _validate_hash_algorithm_support(self) -> None:
        """Validate that the configured hash algorithm is supported by Ibis backend.

        Raises:
            ValueError: If hash algorithm is not supported
        """
        # Create hash functions to check what's supported
        hash_functions = self._create_hash_functions()

        if self.hash_algorithm not in hash_functions:
            supported = [algo.value for algo in hash_functions.keys()]
            raise HashAlgorithmNotSupportedError(
                f"Hash algorithm '{self.hash_algorithm.value}' not supported. "
                f"Supported algorithms: {', '.join(supported)}"
            )

    @property
    def conn(self) -> "SQLBackend":
        """Get Ibis backend connection.

        Returns:
            Active Ibis backend connection

        Raises:
            StoreNotOpenError: If store is not open
        """

        if self._conn is None:
            raise StoreNotOpenError("Ibis connection is not open. Store must be used as a context manager.")
        else:
            return self._conn

    def _open(self, mode: AccessMode) -> None:  # noqa: ARG002
        import ibis

        if self.connection_string:
            self._conn = ibis.connect(self.connection_string)  # ty: ignore[invalid-assignment]
        else:
            assert self.backend is not None, "backend must be set if connection_string is None"
            backend_module = getattr(ibis, self.backend)
            self._conn = backend_module.connect(**self.connection_params)

    def _close(self) -> None:
        if self._conn is not None:
            self._conn.disconnect()
        self._conn = None

    @property
    def sqlalchemy_url(self) -> str:
        """Get SQLAlchemy-compatible connection URL for tools like Alembic.

        Returns the connection string if available. If the store was initialized
        with backend + connection_params instead of a connection string, raises
        an error since constructing a proper URL is backend-specific.

        Returns:
            SQLAlchemy-compatible URL string

        Raises:
            ValueError: If connection_string is not available

        Example:
            <!-- skip next -->
            ```python
            store = IbisMetadataStore("postgresql://user:pass@host:5432/db")
            print(store.sqlalchemy_url)  # postgresql://user:pass@host:5432/db
            ```
        """
        if self.connection_string:
            return self.connection_string

        raise ValueError(
            "SQLAlchemy URL not available. Store was initialized with backend + connection_params "
            "instead of a connection string. To use Alembic, initialize with a connection string: "
            f"IbisMetadataStore('postgresql://user:pass@host:5432/db') instead of "
            f"IbisMetadataStore(backend='{self.backend}', connection_params={{...}})"
        )

    def _write_feature(
        self,
        feature_key: FeatureKey,
        df: Frame,
        **kwargs: Any,
    ) -> None:
        """
        Internal write implementation using Ibis.

        Args:
            feature_key: Feature key to write to
            df: DataFrame with metadata (already validated)
            **kwargs: Backend-specific parameters (currently unused)

        Raises:
            TableNotFoundError: If table doesn't exist and auto_create_tables is False
        """
        table_name = self.get_table_name(feature_key)

        # Apply backend-specific transformations before writing
        df = self.transform_before_write(df, feature_key, table_name)

        from metaxy.config import MetaxyConfig

        if MetaxyConfig.get().enable_map_datatype:
            df = self._handle_map_columns(df, table_name)

        if df.implementation == nw.Implementation.IBIS:
            df_to_insert = df.to_native()  # Ibis expression
        else:
            df_to_insert = (df.collect() if isinstance(df, nw.LazyFrame) else df).to_arrow()

        try:
            self.conn.insert(table_name, obj=df_to_insert)
        except Exception as e:
            if not self._is_table_not_found_error(e):
                raise
            if self.auto_create_tables:
                # Warn about auto-create (first time only)
                if self._should_warn_auto_create_tables:
                    import warnings

                    warnings.warn(
                        f"AUTO_CREATE_TABLES is enabled - automatically creating table '{table_name}'. "
                        "Do not use in production! "
                        "Use proper database migration tools like Alembic for production deployments.",
                        UserWarning,
                        stacklevel=4,
                    )

                # Note: create_table(table_name, obj=df) both creates the table AND inserts the data
                # No separate insert needed - the data from df is already written
                self.conn.create_table(table_name, obj=df_to_insert)
            else:
                raise TableNotFoundError(
                    f"Table '{table_name}' does not exist for feature {feature_key.to_string()}. "
                    f"Enable auto_create_tables=True to automatically create tables, "
                    f"or use proper database migration tools like Alembic to create the table first."
                ) from e

    def _drop_feature(self, feature_key: FeatureKey) -> None:
        """Drop the table for a feature.

        Args:
            feature_key: Feature key to drop metadata for
        """
        table_name = self.get_table_name(feature_key)

        # Check if table exists
        if table_name in self.conn.list_tables():
            self.conn.drop_table(table_name)

    def _get_filtered_ibis_lazy(
        self,
        feature: CoercibleToFeatureKey,
        *,
        filters: Sequence[nw.Expr] | None = None,
    ) -> nw.LazyFrame[Any] | None:
        """Get an Ibis-backed Narwhals LazyFrame with filters applied.

        Always returns Ibis-backed frames regardless of subclass overrides.

        Args:
            feature: Feature to read
            filters: Narwhals filter expressions (converted to WHERE clauses)

        Returns:
            Ibis-backed Narwhals LazyFrame, or None if table doesn't exist
        """
        feature_key = self._resolve_feature_key(feature)
        table_name = self.get_table_name(feature_key)

        import ibis.common.exceptions

        try:
            table = self.conn.table(table_name)
        except ibis.common.exceptions.TableNotFound:
            return None

        table = self.transform_after_read(table, feature_key)

        nw_frame = nw.from_native(table, eager_only=False)
        if not isinstance(nw_frame, nw.LazyFrame):
            raise TypeError(f"Expected narwhals LazyFrame from Ibis table, got {type(nw_frame)}")
        nw_lazy: nw.LazyFrame[Any] = nw_frame

        if filters:
            nw_lazy = nw_lazy.filter(*filters)

        return nw_lazy

    def _delete_feature(
        self,
        feature_key: FeatureKey,
        filters: Sequence[nw.Expr] | None,
        *,
        with_feature_history: bool,  # noqa: ARG002 - version filtering handled by base class
    ) -> None:
        """Backend-specific hard delete implementation for SQL databases.

        Translates Narwhals filter expression to Ibis and executes DELETE in database.

        Args:
            feature_key: Feature to delete from
            filters: Narwhals expressions to filter records
            with_feature_history: Not used here - version filtering handled by base class
        """
        from metaxy.metadata_store.utils import (
            _extract_where_expression,
            _strip_table_qualifiers,
        )

        table_name = self.get_table_name(feature_key)
        filter_list = list(filters or [])

        # Handle empty filters - truncate entire table
        if not filter_list:
            if table_name not in self.conn.list_tables():
                raise TableNotFoundError(f"Table '{table_name}' does not exist for feature {feature_key.to_string()}.")
            self.conn.truncate_table(table_name)
            return

        nw_lazy = self._get_filtered_ibis_lazy(feature_key, filters=filter_list or None)
        if nw_lazy is None:
            raise FeatureNotFoundError(f"Feature {feature_key.to_string()} not found in store")

        # Extract WHERE clause from compiled SELECT statement
        import ibis

        ibis_filtered = nw_lazy.to_native()
        if not isinstance(ibis_filtered, ibis.expr.types.Table):
            raise TypeError(f"Expected nw_lazy.to_native() to return an Ibis Table, got {type(ibis_filtered)!r}")
        select_sql = str(ibis_filtered.compile())

        dialect = self._sql_dialect
        predicate = _extract_where_expression(select_sql, dialect=dialect)
        if predicate is None:
            raise ValueError(f"Cannot extract WHERE clause for DELETE on {self.__class__.__name__}")

        predicate = predicate.transform(_strip_table_qualifiers())
        where_clause = predicate.sql(dialect=dialect) if dialect else predicate.sql()

        delete_stmt = f"DELETE FROM {table_name} WHERE {where_clause}"
        self.conn.raw_sql(delete_stmt)  # ty: ignore[unresolved-attribute]

    @property
    def _sql_dialect(self) -> str | None:
        """Extract SQL dialect from the active backend connection."""
        return self.conn.name

    def _read_feature(
        self,
        feature: CoercibleToFeatureKey,
        *,
        feature_version: str | None = None,
        filters: Sequence[nw.Expr] | None = None,
        columns: Sequence[str] | None = None,
        **kwargs: Any,
    ) -> nw.LazyFrame[Any] | None:
        """
        Read metadata from this store only (no fallback).

        Args:
            feature: Feature to read
            feature_version: Filter by specific feature_version (applied as SQL WHERE clause)
            filters: List of Narwhals filter expressions (converted to SQL WHERE clauses)
            columns: Optional list of columns to select
            **kwargs: Backend-specific parameters (currently unused)

        Returns:
            Narwhals LazyFrame with metadata, or None if not found
        """
        all_filters = list(filters or [])
        if feature_version is not None:
            all_filters.append(nw.col("metaxy_feature_version") == feature_version)

        nw_lazy = self._get_filtered_ibis_lazy(feature, filters=all_filters or None)
        if nw_lazy is None:
            return None

        if columns is not None:
            nw_lazy = nw_lazy.select(columns)

        return nw_lazy

    def ibis_type_to_polars(self, ibis_type: Any) -> Any:
        """Convert an Ibis data type to the corresponding Polars data type.

        Delegates to Ibis's built-in ``DataType.to_polars()`` for most types.
        Provides fallbacks for types that Ibis cannot convert natively:

        - ``Float16`` → ``pl.Float32``
        - ``UUID`` → ``pl.String``
        - ``JSON`` → ``pl.String``
        - ``MACADDR`` → ``pl.String``
        - ``INET`` → ``pl.String``
        - ``GeoSpatial`` → ``pl.Binary``

        Subclasses can override this method to handle store-specific types.

        Args:
            ibis_type: Ibis data type instance

        Returns:
            Corresponding Polars data type

        Raises:
            NotImplementedError: For types with no reasonable Polars mapping
        """
        import ibis.expr.datatypes as dt
        import polars as pl

        try:
            return ibis_type.to_polars()
        except NotImplementedError:
            # Handle types that Ibis's to_polars() doesn't support
            if isinstance(ibis_type, dt.Float16):
                return pl.Float32
            elif isinstance(ibis_type, dt.UUID):
                return pl.String
            elif isinstance(ibis_type, dt.JSON):
                return pl.String
            elif isinstance(ibis_type, dt.MACADDR):
                return pl.String
            elif isinstance(ibis_type, dt.INET):
                return pl.String
            elif isinstance(ibis_type, dt.GeoSpatial):
                return pl.Binary
            else:
                raise

    def transform_after_read(self, table: "ibis.Table", feature_key: "FeatureKey") -> "ibis.Table":
        """Transform Ibis table before wrapping with Narwhals.

        Override in subclasses to apply backend-specific transformations.

        !!! example
            ClickHouse needs to cast JSON columns to Struct for
            PyArrow compatibility.

        Args:
            table: Ibis table reference
            feature_key: The feature key being read (use to get field names)

        Returns:
            Transformed Ibis table (default: unchanged)
        """
        return table

    def transform_before_write(self, df: Frame, feature_key: "FeatureKey", table_name: str) -> Frame:
        """Transform DataFrame before writing to the store.

        Override in subclasses to apply backend-specific transformations.

        Args:
            df: Narwhals DataFrame to be written
            feature_key: The feature key being written to
            table_name: The target table name

        Returns:
            Transformed DataFrame (default: unchanged)
        """
        return df

    @staticmethod
    def _is_table_not_found_error(e: Exception) -> bool:
        """Check if an exception indicates a missing table.

        Override in subclasses to handle backend-specific exception types.
        """
        import ibis.common.exceptions

        return isinstance(e, ibis.common.exceptions.TableNotFound)

    def _handle_map_columns(self, df: Frame, table_name: str) -> Frame:
        """Convert Struct columns to Map format for metaxy system columns.

        For Ibis-backed frames, the conversion stays lazy using ibis.map() expressions.
        For non-Ibis frames, converts through Arrow using polars-map utilities.
        """
        import ibis.expr.datatypes as dt

        from metaxy.models.constants import METAXY_DATA_VERSION_BY_FIELD, METAXY_PROVENANCE_BY_FIELD

        # Always convert the two metaxy system struct columns
        map_columns: set[str] = {METAXY_PROVENANCE_BY_FIELD, METAXY_DATA_VERSION_BY_FIELD}

        # Also convert any user-defined Map columns if table already exists with Map schema
        table_schema: dict[str, dt.DataType] = {}
        if table_name in self.conn.list_tables():
            schema = self.conn.table(table_name).schema()
            table_schema = {name: dtype for name, dtype in schema.items() if isinstance(dtype, dt.Map)}
            map_columns |= set(table_schema)

        if df.implementation == nw.Implementation.IBIS:
            return self._handle_ibis_map_columns(df, map_columns, table_schema)

        return self._handle_polars_map_columns(df, map_columns)

    def _handle_ibis_map_columns(
        self, df: Frame, map_columns: set[str], table_schema: dict[str, "ibis.expr.datatypes.DataType"]
    ) -> Frame:
        """Convert Ibis Struct columns to Map(String, value_type) expressions (stays lazy).

        Keys are always strings (field names from versioning/provenance columns).
        """
        from typing import cast as typing_cast

        import ibis
        import ibis.expr.datatypes as dt

        ibis_table = typing_cast("ibis.Table", df.to_native())
        schema = ibis_table.schema()

        mutations: dict[str, ibis.Expr] = {}
        for col_name in map_columns:
            if col_name not in schema:
                continue
            col_dtype = schema[col_name]
            if not isinstance(col_dtype, dt.Struct):
                continue

            field_names = list(col_dtype.names)  # ty: ignore[invalid-argument-type]

            # Use target Map value type from table schema if available, else infer from struct
            if col_name in table_schema:
                value_type = table_schema[col_name].value_type  # ty: ignore[unresolved-attribute]
            else:
                value_type = col_dtype.types[0] if col_dtype.types else dt.String()  # ty: ignore[not-subscriptable]

            if not field_names:
                mutations[col_name] = ibis.literal(
                    {},
                    type=dt.Map(dt.String(), value_type),  # ty: ignore[invalid-argument-type, missing-argument]
                )
                continue

            keys = ibis.array([ibis.literal(name) for name in field_names])
            values = ibis.array([ibis_table[col_name][name].cast(value_type) for name in field_names])
            mutations[col_name] = ibis.map(keys, values)

        if not mutations:
            return df

        return nw.from_native(ibis_table.mutate(**mutations), eager_only=False)  # ty: ignore[invalid-argument-type]

    def _handle_polars_map_columns(self, df: Frame, map_columns: set[str]) -> Frame:
        """Convert Polars Struct columns to native Arrow Map columns."""
        from metaxy._utils import collect_to_polars
        from metaxy.versioning._arrow_map import convert_extension_maps_to_native, convert_structs_to_maps

        polars_df = collect_to_polars(df)
        polars_df = convert_structs_to_maps(polars_df, columns=list(map_columns))
        return nw.from_native(convert_extension_maps_to_native(polars_df.to_arrow()), eager_only=False)

    def _can_compute_native(self) -> bool:
        """
        Ibis backends support native field provenance calculations (Narwhals-based).

        Returns:
            True (use Narwhals components with Ibis-backed tables)

        Note: All Ibis stores now use Narwhals-based components (NarwhalsJoiner,
        PolarsProvenanceByFieldCalculator, NarwhalsDiffResolver) which work efficiently
        with Ibis-backed tables.
        """
        return True

    def display(self) -> str:
        """Display string for this store."""
        from metaxy.metadata_store.utils import sanitize_uri

        backend_info = self.connection_string or f"{self.backend}"
        # Sanitize connection strings that may contain credentials
        sanitized_info = sanitize_uri(backend_info)
        return f"{self.__class__.__name__}(backend={sanitized_info})"

    def _get_store_metadata_impl(self, feature_key: CoercibleToFeatureKey) -> dict[str, Any]:
        """Return store metadata including table name.

        Args:
            feature_key: Feature key to get metadata for.

        Returns:
            Dictionary with `table_name` key.
        """
        resolved_key = self._resolve_feature_key(feature_key)
        return {"table_name": self.get_table_name(resolved_key)}

    @classmethod
    def config_model(cls) -> type[IbisMetadataStoreConfig]:
        return IbisMetadataStoreConfig
