"""
Table and Column convenience wrappers for the Eddytor SDK.

Provides a fluent API so callers don't have to repeat catalog/schema/table
on every call:

    table = client.table("eddytor", "cfg_xxx", "demo_products")
    print(table.count())
    print(table.history())
    table.column("status").set_fixed_domain(["Active", "Discontinued"])
"""

from __future__ import annotations

import warnings
from typing import TYPE_CHECKING, Any, Union

import pyarrow as pa
import pyarrow.compute as pc

from .exceptions import EddytorSchemaError

if TYPE_CHECKING:
    import pandas as pd

    from .client import EddytorClient
    from .grpc_client import StreamResult
    from .models import (
        AllowedValues,
        ColumnDomain,
        ExplainRowsResult,
        MagicDustResult,
        MoveTableResult,
        TableHistory,
        TableMetadata,
    )

_VALID_MERGE_OPS = {"INSERT", "UPDATE", "DELETE"}


class Table:
    """
    Lightweight handle for a ``catalog.schema.table`` triple.

    All methods delegate to the parent :class:`EddytorClient` or its
    :class:`RestClient`, filling in the three-part identifier automatically.
    """

    def __init__(self, client: EddytorClient, catalog: str, schema: str, table: str) -> None:
        self._client = client
        self._catalog = catalog
        self._schema = schema
        self._table = table
        self._cached_schema: pa.Schema | None = None
        self._cached_pks: list[str] | None = None

    @property
    def catalog(self) -> str:
        return self._catalog

    @property
    def schema(self) -> str:
        return self._schema

    @property
    def name(self) -> str:
        return self._table

    @property
    def fqn(self) -> str:
        """Fully-qualified name: ``catalog.schema.name``."""
        return f"`{self._catalog}`.`{self._schema}`.`{self._table}`"

    def __repr__(self) -> str:
        return f"Table({self.fqn!r})"

    # ------------------------------------------------------------------
    # Schema caching
    # ------------------------------------------------------------------

    def _fetch_schema(self) -> pa.Schema:
        """Fetch the table's Arrow schema via a LIMIT 0 query."""
        return self._client.query_arrow(f"SELECT * FROM {self.fqn} LIMIT 0").schema

    def _get_schema(self) -> pa.Schema:
        """Get cached schema, fetching on first call."""
        if self._cached_schema is None:
            self._cached_schema = self._fetch_schema()
        return self._cached_schema

    def _get_primary_keys(self) -> list[str]:
        """Extract PK column names from schema field metadata.

        Tries Arrow field metadata first (``isPrimaryKey: "true"``).
        Falls back to the REST ``get_table_metadata()`` endpoint if no
        PKs are found in the Arrow metadata (DataFusion may strip field
        metadata during query planning).
        """
        if self._cached_pks is None:
            schema = self._get_schema()
            pks = [
                schema.field(i).name
                for i in range(len(schema))
                if schema.field(i).metadata
                and schema.field(i).metadata.get(b"isPrimaryKey") == b"true"
            ]
            if not pks:
                try:
                    meta = self._client.rest.get_table_metadata(
                        self._catalog, self._schema, self._table,
                    )
                    pks = [
                        col.name
                        for col in meta.columns
                        if col.metadata and col.metadata.get("isPrimaryKey") == "true"
                    ]
                except Exception:
                    pks = []
            self._cached_pks = pks
        return self._cached_pks

    def refresh_schema(self) -> None:
        """Force-refresh the cached schema (call after DDL operations)."""
        self._cached_schema = None
        self._cached_pks = None

    # ------------------------------------------------------------------
    # Data alignment
    # ------------------------------------------------------------------

    @staticmethod
    def _to_arrow_table(data: Union[pa.Table, pa.RecordBatch, "pd.DataFrame"]) -> pa.Table:
        """Normalize input data to a :class:`pa.Table`."""
        if isinstance(data, pa.Table):
            return data
        if isinstance(data, pa.RecordBatch):
            return pa.Table.from_batches([data])
        # pandas DataFrame
        return pa.Table.from_pandas(data)

    def _prepare_data(
        self,
        data: Union[pa.Table, pa.RecordBatch, "pd.DataFrame"],
        operation: str,
    ) -> Union[pa.Table, pa.RecordBatch, "pd.DataFrame"]:
        """Align *data* to the table's schema for the given *operation*.

        Returns the same container type as the input (Table, RecordBatch,
        or DataFrame).
        """
        was_batch = isinstance(data, pa.RecordBatch)
        was_df = not isinstance(data, (pa.Table, pa.RecordBatch))
        table = self._to_arrow_table(data)

        if operation == "insert":
            aligned = self._align_for_insert(table)
        elif operation == "merge":
            aligned = self._align_for_merge(table)
        elif operation == "delete":
            aligned = self._align_for_delete(table)
        else:
            raise ValueError(f"Unknown operation: {operation!r}")

        if was_batch:
            return aligned.to_batches()[0] if aligned.num_rows else pa.RecordBatch.from_pydict(
                {name: [] for name in aligned.schema.names}, schema=aligned.schema,
            )
        if was_df:
            return aligned.to_pandas()
        return aligned

    def _align_for_insert(self, data: pa.Table) -> pa.Table:
        """Align *data* for an ``INSERT`` operation.

        1. Drop extra columns (warns).
        2. Add missing *nullable* columns as null arrays (warns).
        3. Error on missing *non-nullable* columns.
        4. Cast column types (safe cast).
        5. Reorder to match the target schema.
        """
        target = self._get_schema()
        target_names = set(target.names)
        data_names = set(data.schema.names)

        # --- extra columns ---
        extras = data_names - target_names
        if extras:
            warnings.warn(
                f"Dropping columns not in target schema: {sorted(extras)}",
                stacklevel=3,
            )
            data = data.drop_columns(sorted(extras))

        # --- missing / cast ---
        columns: list[pa.ChunkedArray] = []
        fields: list[pa.Field] = []
        for i in range(len(target)):
            target_field = target.field(i)
            col_name = target_field.name

            if col_name in data.schema.names:
                col = data.column(col_name)
                if col.type != target_field.type:
                    try:
                        col = pc.cast(col, target_field.type, safe=True)
                    except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as exc:
                        raise EddytorSchemaError(
                            f"Cannot cast column {col_name!r} from {col.type} "
                            f"to {target_field.type}: {exc}"
                        ) from exc
                columns.append(col)
                fields.append(target_field)
            elif target_field.nullable:
                warnings.warn(
                    f"Adding missing nullable column {col_name!r} as nulls",
                    stacklevel=3,
                )
                null_array = pa.nulls(data.num_rows, type=target_field.type)
                columns.append(null_array)
                fields.append(target_field)
            else:
                raise EddytorSchemaError(
                    f"Missing required non-nullable column: {col_name!r}"
                )

        return pa.table(
            {field.name: col for field, col in zip(fields, columns)},
            schema=pa.schema(fields),
        )

    def _align_for_merge(self, data: pa.Table) -> pa.Table:
        """Align *data* for a ``MERGE`` operation.

        Validates the ``_operation`` column, strips it, aligns the
        remaining columns via insert alignment, then re-attaches
        ``_operation``.
        """
        if "_operation" not in data.schema.names:
            raise EddytorSchemaError(
                "Merge data must contain an '_operation' column "
                "(values: 'INSERT', 'UPDATE', 'DELETE')"
            )

        op_col = data.column("_operation")
        if not pa.types.is_string(op_col.type) and not pa.types.is_large_string(op_col.type):
            raise EddytorSchemaError(
                f"'_operation' column must be a string type, got {op_col.type}"
            )

        unique_ops = pc.unique(op_col).to_pylist()
        invalid = set(unique_ops) - _VALID_MERGE_OPS
        if invalid:
            raise EddytorSchemaError(
                f"Invalid '_operation' values: {sorted(invalid)}. "
                f"Allowed: {sorted(_VALID_MERGE_OPS)}"
            )

        # Strip _operation, align remaining columns, then re-attach
        remaining = data.drop_columns(["_operation"])
        aligned = self._align_for_insert(remaining)
        return aligned.append_column(
            pa.field("_operation", pa.string()),
            op_col,
        )

    def _align_for_delete(self, data: pa.Table) -> pa.Table:
        """Align *data* for a ``DELETE`` operation.

        1. Validate all provided columns exist in the target schema.
        2. If PKs exist, require all PK columns present.
        3. Cast types to match target.
        """
        target = self._get_schema()
        target_names = set(target.names)
        data_names = set(data.schema.names)

        unknown = data_names - target_names
        if unknown:
            raise EddytorSchemaError(
                f"Delete data contains columns not in the table schema: {sorted(unknown)}"
            )

        pks = self._get_primary_keys()
        if pks:
            missing_pks = set(pks) - data_names
            if missing_pks:
                raise EddytorSchemaError(
                    f"Delete data is missing primary key columns: {sorted(missing_pks)}"
                )

        # Cast types to match target schema
        columns: list[pa.ChunkedArray] = []
        fields: list[pa.Field] = []
        for col_name in data.schema.names:
            target_field = target.field(target.get_field_index(col_name))
            col = data.column(col_name)
            if col.type != target_field.type:
                try:
                    col = pc.cast(col, target_field.type, safe=True)
                except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as exc:
                    raise EddytorSchemaError(
                        f"Cannot cast column {col_name!r} from {col.type} "
                        f"to {target_field.type}: {exc}"
                    ) from exc
            columns.append(col)
            fields.append(target_field.with_name(col_name))

        return pa.table(
            {field.name: col for field, col in zip(fields, columns)},
            schema=pa.schema(fields),
        )

    # ------------------------------------------------------------------
    # Flight SQL helpers
    # ------------------------------------------------------------------

    def query(self, sql: str) -> pd.DataFrame:
        """Execute a SQL query and return a pandas DataFrame."""
        return self._client.query(sql)

    def query_all(self, limit: int | None = None) -> pd.DataFrame:
        """``SELECT * FROM <fqn> [LIMIT n]`` as a pandas DataFrame."""
        sql = f"SELECT * FROM {self.fqn}"
        if limit is not None:
            sql += f" LIMIT {int(limit)}"
        return self._client.query(sql)

    def query_arrow(self, sql: str) -> pa.Table:
        """Execute a SQL query and return a PyArrow Table."""
        return self._client.query_arrow(sql)

    def query_all_arrow(self, limit: int | None = None) -> pa.Table:
        """``SELECT * FROM <fqn> [LIMIT n]`` as a PyArrow Table."""
        sql = f"SELECT * FROM {self.fqn}"
        if limit is not None:
            sql += f" LIMIT {int(limit)}"
        return self._client.query_arrow(sql)

    def execute(self, sql: str) -> list[tuple[Any, ...]]:
        """Execute a SQL query and return a list of tuples."""
        return self._client.execute(sql)

    def execute_dml(self, sql: str) -> int:
        """Execute a DML statement and return the affected row count."""
        return self._client.execute_dml(sql)

    def describe(self) -> pd.DataFrame:
        """Get the schema/columns of this table."""
        return self._client.describe_table(self.fqn)

    def exists(self) -> bool:
        """Check if this table exists."""
        return self._client.table_exists(self.fqn)

    def count(self, where: str | None = None) -> int:
        """Count rows, with an optional WHERE clause (without the ``WHERE`` keyword)."""
        return self._client.count(self.fqn, where)

    def ingest(self, data: Any, mode: str = "append") -> int:
        """Bulk ingest Arrow data into this table via ADBC ingest."""
        return self._client.ingest(
            self._table, data, mode=mode,
            catalog=self._catalog, schema=self._schema,
        )

    def add_column(
        self,
        columns: pa.Schema,
        constraints: list[dict[str, str]] | None = None,
    ) -> None:
        """Add columns to this table via Flight SQL DoAction."""
        self._client.add_column(
            self._table, columns,
            catalog=self._catalog, schema=self._schema,
            constraints=constraints,
        )
        self.refresh_schema()

    # ------------------------------------------------------------------
    # Native gRPC data operations (the "easy" API)
    # ------------------------------------------------------------------

    def stream(
        self,
        *,
        limit: int | None = None,
        offset: int | None = None,
        filters: list[dict[str, Any]] | None = None,
        search: str | None = None,
        search_columns: list[str] | None = None,
        sort: list[dict[str, Any]] | None = None,
    ) -> StreamResult:
        """Stream data from this table with optional filtering, sorting, and pagination.

        Uses the native EddytorGRPC service for structured queries.

        Args:
            limit: Maximum rows to return.
            offset: Number of rows to skip.
            filters: List of filter dicts with ``column``, ``operator``, and
                     ``values`` keys.
            search: Free-text search term.
            search_columns: Columns to search (default: all string columns).
            sort: List of sort dicts with ``column`` and optionally
                  ``descending`` (bool) keys.

        Returns:
            StreamResult wrapping a PyArrow Table and optional total_rows.

        Example:
            >>> result = table.stream(limit=10, sort=[{"column": "name"}])
            >>> print(result.total_rows)
            >>> print(result.to_pandas())
        """
        return self._client.grpc.stream_data(
            self._catalog, self._schema, self._table,
            limit=limit, offset=offset, filters=filters,
            search=search, search_columns=search_columns, sort=sort,
        )

    def insert(
        self,
        data: Union[pa.Table, pa.RecordBatch, "pd.DataFrame"],
        *,
        auto_align: bool = True,
    ) -> None:
        """Insert rows into this table via native gRPC.

        Args:
            data: Data to insert (DataFrame, Arrow Table, or RecordBatch).
            auto_align: If ``True`` (default), automatically align data to the
                table's schema — dropping extra columns, adding missing nullable
                columns as nulls, and casting types.

        Example:
            >>> import pyarrow as pa
            >>> data = pa.table({"id": [1, 2], "name": ["a", "b"]})
            >>> table.insert(data)
        """
        if auto_align:
            data = self._prepare_data(data, "insert")
        self._client.grpc.insert_into_table(
            self._catalog, self._schema, self._table, data,
        )

    def merge(
        self,
        data: Union[pa.Table, pa.RecordBatch, "pd.DataFrame"],
        comment: str | None = None,
        *,
        auto_align: bool = True,
    ) -> None:
        """Merge (upsert) rows in this table via native gRPC.

        Args:
            data: Data to merge (DataFrame, Arrow Table, or RecordBatch).
            comment: Optional commit comment.
            auto_align: If ``True`` (default), automatically align data to the
                table's schema — validating the ``_operation`` column, dropping
                extra columns, adding missing nullable columns, and casting types.

        Example:
            >>> data = pa.table({"id": [1], "name": ["updated"]})
            >>> table.merge(data, comment="Update name for id 1")
        """
        if auto_align:
            data = self._prepare_data(data, "merge")
        self._client.grpc.merge_table(
            self._catalog, self._schema, self._table, data, comment=comment,
        )

    def delete(
        self,
        data: Union[pa.Table, pa.RecordBatch, "pd.DataFrame"],
        *,
        auto_align: bool = True,
    ) -> None:
        """Delete rows from this table via native gRPC.

        Args:
            data: Data identifying rows to delete (DataFrame, Arrow Table, or RecordBatch).
            auto_align: If ``True`` (default), validate that all columns exist
                in the table schema, ensure PK columns are present, and cast types.

        Example:
            >>> data = pa.table({"id": [1, 2]})
            >>> table.delete(data)
        """
        if auto_align:
            data = self._prepare_data(data, "delete")
        self._client.grpc.delete_from_table(
            self._catalog, self._schema, self._table, data,
        )

    # ------------------------------------------------------------------
    # REST helpers
    # ------------------------------------------------------------------

    def metadata(self) -> TableMetadata:
        """Get table metadata (columns, partitions, etc.)."""
        return self._client.rest.get_table_metadata(self._catalog, self._schema, self._table)

    def history(self) -> TableHistory:
        """Get the version history of this table."""
        return self._client.rest.get_table_history(self._catalog, self._schema, self._table)

    def rollback(self, version: int) -> None:
        """Rollback this table to a specific version."""
        self._client.rest.rollback_table(
            self._catalog, self._schema, self._table, version)

    def add_constraints(self, constraints: list[dict[str, Any]]) -> None:
        """Add constraints to this table."""
        self._client.rest.add_constraints(
            self._catalog, self._schema, self._table, constraints)

    def drop_constraint(self, name: str) -> None:
        """Drop a constraint from this table."""
        self._client.rest.drop_constraint(
            self._catalog, self._schema, self._table, name)

    def update_field_metadata(self, field_table: str, metadata: dict[str, Any]) -> None:
        """Update metadata for a table field."""
        self._client.rest.update_field_metadata(
            self._catalog, self._schema, self._table, field_table, metadata)

    def check_enum_value_usage(self, column: str, value: str) -> bool:
        """Check if an enum value is in use (for safe deletion)."""
        return self._client.rest.check_enum_value_usage(self._catalog, self._schema, self._table, column, value)

    def magic_dust(
        self,
        provider: str,
        action: str,
        model: str,
        max_tokens: int = 4096,
        sample_size: int = 1000,
    ) -> MagicDustResult:
        """Execute AI-powered analysis on this table's data."""
        return self._client.rest.magic_dust(
            self._catalog, self._schema, self._table,
            provider=provider, action=action, model=model,
            max_tokens=max_tokens, sample_size=sample_size,
        )

    def explain_rows(
        self,
        provider: str,
        model: str,
        row_pks: list[Any] | None = None,
        pk_column: str | None = None,
        sample_size: int = 20,
        max_tokens: int = 4096,
    ) -> ExplainRowsResult:
        """Generate a narrative explanation of selected rows with tagged cell references."""
        return self._client.rest.explain_rows(
            self._catalog, self._schema, self._table,
            provider=provider, model=model,
            row_pks=row_pks, pk_column=pk_column,
            sample_size=sample_size, max_tokens=max_tokens,
        )

    def drop(self) -> None:
        """Permanently delete this table and all its data from the object store.

        This action is irreversible.
        """
        self._client.rest.delete_table(self._catalog, self._schema, self._table)

    def move(self, destination_config_id: str, destination_path: str) -> MoveTableResult:
        """Move this table to a different storage configuration.

        The table is deregistered from the source and re-discovered at
        the destination.
        """
        return self._client.rest.move_table(
            self._catalog, self._schema, self._table,
            destination_config_id, destination_path,
        )

    # ------------------------------------------------------------------
    # Column accessor
    # ------------------------------------------------------------------

    def column(self, name: str) -> Column:
        """Get a :class:`Column` handle for the given column name."""
        return Column(self, name)


class Column:
    """
    Lightweight handle for a single column within a :class:`Table`.

    All methods delegate to the parent table's REST client, filling in
    catalog, schema, table, and column automatically.
    """

    def __init__(self, table: Table, name: str) -> None:
        self._table = table
        self._name = name

    @property
    def name(self) -> str:
        return self._name

    def __repr__(self) -> str:
        return f"Column({self._table.fqn}.{self._name!r})"

    @property
    def _rest(self):
        return self._table._client.rest

    @property
    def _c(self) -> str:
        return self._table._catalog

    @property
    def _s(self) -> str:
        return self._table._schema

    @property
    def _t(self) -> str:
        return self._table._table

    # ------------------------------------------------------------------
    # Domain CRUD
    # ------------------------------------------------------------------

    def domain(self) -> ColumnDomain:
        """Get the domain configuration for this column."""
        return self._rest.get_column_domain(self._c, self._s, self._t, self._name)

    def set_fixed_domain(self, values: list[str]) -> ColumnDomain:
        """Set a fixed domain (enum) on this column."""
        return self._rest.set_fixed_domain(self._c, self._s, self._t, self._name, values)

    def set_hierarchical_inline_domain(self, hierarchy: dict[str, list[str]]) -> ColumnDomain:
        """Set a hierarchical inline domain on this column."""
        return self._rest.set_hierarchical_inline_domain(self._c, self._s, self._t, self._name, hierarchy)

    def set_hierarchical_derived_domain(self, parent_column: str) -> ColumnDomain:
        """Set a hierarchical derived domain (linked to another column)."""
        return self._rest.set_hierarchical_derived_domain(self._c, self._s, self._t, self._name, parent_column)

    def set_hierarchical_table_domain(
        self,
        source_table: str,
        source_column: str,
        parent_column: str | None = None,
    ) -> ColumnDomain:
        """Set a hierarchical domain from another table."""
        return self._rest.set_hierarchical_table_domain(
            self._c, self._s, self._t, self._name,
            source_table, source_column, parent_column,
        )

    def delete_domain(self) -> None:
        """Remove the domain from this column."""
        self._rest.delete_column_domain(self._c, self._s, self._t, self._name)

    # ------------------------------------------------------------------
    # Fixed domain value operations
    # ------------------------------------------------------------------

    def add_fixed_value(self, value: str) -> ColumnDomain:
        """Add a value to the fixed domain."""
        return self._rest.add_fixed_value(self._c, self._s, self._t, self._name, value)

    def set_fixed_values(self, values: list[str]) -> ColumnDomain:
        """Replace all values in the fixed domain."""
        return self._rest.set_fixed_values(self._c, self._s, self._t, self._name, values)

    def remove_fixed_value(self, value: str) -> ColumnDomain:
        """Remove a value from the fixed domain."""
        return self._rest.remove_fixed_value(self._c, self._s, self._t, self._name, value)

    # ------------------------------------------------------------------
    # Hierarchical domain operations
    # ------------------------------------------------------------------

    def add_hierarchy_parent(self, parent: str) -> ColumnDomain:
        """Add a parent to the hierarchical inline domain."""
        return self._rest.add_hierarchy_parent(self._c, self._s, self._t, self._name, parent)

    def remove_hierarchy_parent(self, parent: str) -> ColumnDomain:
        """Remove a parent from the hierarchical inline domain."""
        return self._rest.remove_hierarchy_parent(self._c, self._s, self._t, self._name, parent)

    def add_hierarchy_child(self, parent: str, child: str) -> ColumnDomain:
        """Add a child to a parent in the hierarchical inline domain."""
        return self._rest.add_hierarchy_child(self._c, self._s, self._t, self._name, parent, child)

    def set_hierarchy_children(self, parent: str, children: list[str]) -> ColumnDomain:
        """Replace all children for a parent in the hierarchical inline domain."""
        return self._rest.set_hierarchy_children(self._c, self._s, self._t, self._name, parent, children)

    def remove_hierarchy_child(self, parent: str, child: str) -> ColumnDomain:
        """Remove a child from a parent in the hierarchical inline domain."""
        return self._rest.remove_hierarchy_child(self._c, self._s, self._t, self._name, parent, child)

    # ------------------------------------------------------------------
    # Allowed values query
    # ------------------------------------------------------------------

    def allowed_values(self, parent_value: str | None = None) -> AllowedValues:
        """Get allowed values for this column, optionally filtered by parent."""
        return self._rest.get_allowed_values(self._c, self._s, self._t, self._name, parent_value)
