"""
Native EddytorGRPC client for structured data operations.

Provides GrpcClient for insert/merge/delete/stream operations using Arrow IPC
over the native gRPC service (port 50051), and StreamResult for paginated query
results.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Union

import grpc
import pyarrow as pa
import pyarrow.ipc as ipc

from .exceptions import EddytorAuthError, EddytorConnectionError, EddytorQueryError
from .proto import EddytorGRPC_pb2, EddytorGRPC_pb2_grpc

if TYPE_CHECKING:
    import pandas as pd

# 64 MB max message size for large batches
_MAX_MESSAGE_LENGTH = 64 * 1024 * 1024

_CHANNEL_OPTIONS = [
    ("grpc.max_send_message_length", _MAX_MESSAGE_LENGTH),
    ("grpc.max_receive_message_length", _MAX_MESSAGE_LENGTH),
]


@dataclass
class StreamResult:
    """Result of a ``stream_data`` call.

    Wraps a :class:`pyarrow.Table` with an optional ``total_rows`` count
    (only available when the server reports it on the first batch).
    """

    table: pa.Table
    total_rows: int | None = None

    def to_pandas(self) -> pd.DataFrame:
        """Convert the result to a pandas DataFrame."""
        return self.table.to_pandas()


class GrpcClient:
    """Client for the native EddytorGRPC service.

    The gRPC channel is created lazily on first use. Authentication is
    Bearer-token based, using the same API key as Flight SQL.
    """

    def __init__(
        self,
        api_key: str,
        host: str = "localhost",
        port: int = 50051,
        use_tls: bool = True,
        skip_verify: bool = False,
    ):
        self._api_key = api_key
        self._host = host
        self._port = port
        self._use_tls = use_tls
        self._skip_verify = skip_verify
        self._channel: grpc.Channel | None = None
        self._stub: EddytorGRPC_pb2_grpc.EddytorGRPCStub | None = None

    # ------------------------------------------------------------------
    # Channel management
    # ------------------------------------------------------------------

    def _ensure_channel(self) -> EddytorGRPC_pb2_grpc.EddytorGRPCStub:
        """Lazily create the gRPC channel and stub."""
        if self._stub is not None:
            return self._stub

        target = f"{self._host}:{self._port}"
        try:
            if self._use_tls:
                if self._skip_verify:
                    # Create a channel with TLS but skip certificate verification
                    credentials = grpc.ssl_channel_credentials()
                    # Override SSL target name to skip verification
                    options = _CHANNEL_OPTIONS + [
                        ("grpc.ssl_target_name_override", self._host),
                    ]
                    self._channel = grpc.secure_channel(target, credentials, options=options)
                else:
                    credentials = grpc.ssl_channel_credentials()
                    self._channel = grpc.secure_channel(target, credentials, options=_CHANNEL_OPTIONS)
            else:
                self._channel = grpc.insecure_channel(target, options=_CHANNEL_OPTIONS)
        except Exception as e:
            raise EddytorConnectionError(f"Failed to create gRPC channel to {target}: {e}") from e

        self._stub = EddytorGRPC_pb2_grpc.EddytorGRPCStub(self._channel)
        return self._stub

    def _metadata(self) -> list[tuple[str, str]]:
        """Build Bearer auth metadata for every RPC call."""
        return [("authorization", f"Bearer {self._api_key}")]

    def close(self) -> None:
        """Close the gRPC channel."""
        if self._channel is not None:
            self._channel.close()
            self._channel = None
            self._stub = None

    # ------------------------------------------------------------------
    # Arrow IPC helpers
    # ------------------------------------------------------------------

    @staticmethod
    def to_ipc(data: Union[pa.Table, pa.RecordBatch, "pd.DataFrame"]) -> bytes:
        """Serialize a DataFrame, Arrow Table, or RecordBatch to Arrow IPC streaming bytes."""
        try:
            import pandas as _pd
            if isinstance(data, _pd.DataFrame):
                data = pa.Table.from_pandas(data)
        except ImportError:
            pass

        if isinstance(data, pa.Table):
            batches = data.to_batches()
            if not batches:
                # Empty table: write a zero-row batch so the schema is present
                empty_arrays = [pa.array([], type=f.type) for f in data.schema]
                batches = [pa.record_batch(empty_arrays, schema=data.schema)]
            schema = data.schema
        elif isinstance(data, pa.RecordBatch):
            batches = [data]
            schema = data.schema
        else:
            raise TypeError(f"Unsupported data type: {type(data)}. Expected DataFrame, pa.Table, or pa.RecordBatch.")

        sink = pa.BufferOutputStream()
        writer = ipc.new_stream(sink, schema)
        for batch in batches:
            writer.write_batch(batch)
        writer.close()
        return sink.getvalue().to_pybytes()

    @staticmethod
    def from_ipc(data: bytes) -> pa.RecordBatch:
        """Deserialize Arrow IPC streaming bytes into a single RecordBatch."""
        reader = ipc.open_stream(data)
        batches = [b for b in reader]
        if not batches:
            return pa.record_batch([], schema=reader.schema)
        if len(batches) == 1:
            return batches[0]
        table = pa.Table.from_batches(batches, schema=reader.schema)
        return table.to_batches()[0] if table.num_rows > 0 else pa.record_batch([], schema=reader.schema)

    # ------------------------------------------------------------------
    # RPC method helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _make_table_ref(
        catalog: str, schema: str, table: str
    ) -> EddytorGRPC_pb2.ProtoTableReference:
        return EddytorGRPC_pb2.ProtoTableReference(
            catalog=catalog, schema=schema, table=table
        )

    def _handle_rpc_error(self, err: grpc.RpcError) -> None:
        """Map a gRPC RpcError to an appropriate SDK exception."""
        code = err.code()
        details = err.details() or str(err)

        if code == grpc.StatusCode.UNAUTHENTICATED:
            raise EddytorAuthError(f"Authentication failed: {details}") from err
        if code == grpc.StatusCode.PERMISSION_DENIED:
            raise EddytorAuthError(f"Permission denied: {details}") from err
        if code == grpc.StatusCode.UNAVAILABLE:
            raise EddytorConnectionError(f"gRPC service unavailable: {details}") from err
        raise EddytorQueryError(f"gRPC error ({code.name}): {details}") from err

    # ------------------------------------------------------------------
    # Data operations
    # ------------------------------------------------------------------

    def stream_data(
        self,
        catalog: str,
        schema: str,
        table: str,
        *,
        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 a table with optional filtering, sorting, and pagination.

        Args:
            catalog: Catalog name.
            schema: Schema name.
            table: Table name.
            limit: Maximum rows to return.
            offset: Number of rows to skip.
            filters: List of filter dicts, each with ``column``, ``operator``,
                     and ``values`` (list of strings) keys.
            search: Free-text search term.
            search_columns: Columns to search in (default: all string columns).
            sort: List of sort dicts, each with ``column`` and optionally
                  ``descending`` (bool, default False) keys.

        Returns:
            StreamResult wrapping a PyArrow Table and optional total_rows.
        """
        stub = self._ensure_channel()

        proto_filters = []
        if filters:
            for f in filters:
                proto_filters.append(EddytorGRPC_pb2.ColumnFilter(
                    column=f["column"],
                    operator=f["operator"],
                    values=f.get("values", []),
                ))

        proto_sort = []
        if sort:
            for s in sort:
                proto_sort.append(EddytorGRPC_pb2.SortColumn(
                    column=s["column"],
                    descending=s.get("descending", False),
                ))

        msg = EddytorGRPC_pb2.QueryMessage(
            table_reference=self._make_table_ref(catalog, schema, table),
            filters=proto_filters,
            search_columns=search_columns or [],
            sort_columns=proto_sort,
        )
        if limit is not None:
            msg.limit = limit
        if offset is not None:
            msg.offset = offset
        if search is not None:
            msg.search_text = search

        try:
            response_iter = stub.StreamData(msg, metadata=self._metadata())
        except grpc.RpcError as e:
            self._handle_rpc_error(e)

        batches: list[pa.RecordBatch] = []
        total_rows: int | None = None
        result_schema: pa.Schema | None = None

        for arrow_batch_msg in response_iter:
            if total_rows is None and arrow_batch_msg.HasField("total_rows"):
                total_rows = arrow_batch_msg.total_rows

            reader = ipc.open_stream(arrow_batch_msg.record_batch)
            if result_schema is None:
                result_schema = reader.schema
            for batch in reader:
                batches.append(batch)

        if result_schema is None:
            result_schema = pa.schema([])

        result_table = pa.Table.from_batches(batches, schema=result_schema) if batches else pa.table({}, schema=result_schema)
        return StreamResult(table=result_table, total_rows=total_rows)

    def insert_into_table(
        self,
        catalog: str,
        schema: str,
        table: str,
        data: Union[pa.Table, pa.RecordBatch, "pd.DataFrame"],
    ) -> None:
        """Insert rows into a table.

        Args:
            catalog: Catalog name.
            schema: Schema name.
            table: Table name.
            data: Data to insert (DataFrame, Arrow Table, or RecordBatch).
        """
        stub = self._ensure_channel()
        ipc_bytes = self.to_ipc(data)

        msg = EddytorGRPC_pb2.TableRecordBatch(
            table_reference=self._make_table_ref(catalog, schema, table),
            record_batch=ipc_bytes,
        )

        try:
            stub.InsertIntoTable(msg, metadata=self._metadata())
        except grpc.RpcError as e:
            self._handle_rpc_error(e)

    def merge_table(
        self,
        catalog: str,
        schema: str,
        table: str,
        data: Union[pa.Table, pa.RecordBatch, "pd.DataFrame"],
        comment: str | None = None,
    ) -> None:
        """Merge (upsert) rows in a table.

        Args:
            catalog: Catalog name.
            schema: Schema name.
            table: Table name.
            data: Data to merge (DataFrame, Arrow Table, or RecordBatch).
            comment: Optional commit comment.
        """
        stub = self._ensure_channel()
        ipc_bytes = self.to_ipc(data)

        msg = EddytorGRPC_pb2.MergeRpcRequest(
            table_reference=self._make_table_ref(catalog, schema, table),
            record_batch=ipc_bytes,
        )
        if comment is not None:
            msg.comment = comment

        try:
            stub.MergeTable(msg, metadata=self._metadata())
        except grpc.RpcError as e:
            self._handle_rpc_error(e)

    def delete_from_table(
        self,
        catalog: str,
        schema: str,
        table: str,
        data: Union[pa.Table, pa.RecordBatch, "pd.DataFrame"],
    ) -> None:
        """Delete rows from a table.

        Args:
            catalog: Catalog name.
            schema: Schema name.
            table: Table name.
            data: Data identifying rows to delete (DataFrame, Arrow Table, or RecordBatch).
        """
        stub = self._ensure_channel()
        ipc_bytes = self.to_ipc(data)

        msg = EddytorGRPC_pb2.TableRecordBatch(
            table_reference=self._make_table_ref(catalog, schema, table),
            record_batch=ipc_bytes,
        )

        try:
            stub.DeleteFromTable(msg, metadata=self._metadata())
        except grpc.RpcError as e:
            self._handle_rpc_error(e)
