"""
Eddytor Flight SQL Client.

Provides a high-level interface to connect to Eddytor via Apache Arrow Flight SQL
and REST APIs for table management, storage, and column domains.
"""

from __future__ import annotations

import json
import os
from typing import TYPE_CHECKING, Any

import adbc_driver_flightsql.dbapi as flight_sql
import pyarrow as pa
import pyarrow.flight as flight

from .exceptions import EddytorConnectionError, EddytorQueryError
from .grpc_client import GrpcClient
from .rest import RestClient
from .storage import Storage
from .table import Table

if TYPE_CHECKING:
    import pandas as pd


class EddytorClient:
    """
    Client for connecting to Eddytor via Flight SQL and REST APIs.

    **Table naming:** Eddytor tables use a three-part Fully Qualified Name (FQN)::

        eddytor.cfg_<config_id_hex>.<hash>_<table_name>

    - **Catalog** is always ``"eddytor"``.
    - **Schema** is ``cfg_`` + storage config UUID (dashes removed).
    - **Table** is a hash prefix + human-readable name.

    Use :meth:`tables` to discover all available tables and get ``Table`` handles
    with the FQN already filled in — you never need to construct these names manually.

    Quick start::

        with EddytorClient(api_key="edd_live_xxx") as client:
            # Discover tables
            for t in client.tables():
                print(t.fqn)

            # Query a table
            df = client.tables()[0].query_all(limit=10)

    Or query with raw SQL (use the FQN from discovery)::

        df = client.query("SELECT * FROM `eddytor`.`cfg_550e...`.`abc_customers` LIMIT 10")
    """

    def __init__(self, api_key: str):
        """
        Initialize the Eddytor client.

        Args:
            api_key: Eddytor API key (starts with 'edd_live_')
        """
        self.api_key = api_key

        # Production endpoints (internal overrides via env vars for development)
        self._flight_host = os.environ.get("EDDYTOR_FLIGHT_HOST", "flight.eddytor.com")
        self._flight_port = int(os.environ.get("EDDYTOR_FLIGHT_PORT", "443"))
        self._api_host = os.environ.get("EDDYTOR_API_HOST", "api.eddytor.com")
        self._api_port = int(os.environ.get("EDDYTOR_API_PORT", "443"))
        self._grpc_host = os.environ.get("EDDYTOR_GRPC_HOST", "api.eddytor.com")
        self._grpc_port = int(os.environ.get("EDDYTOR_GRPC_PORT", "443"))
        self._use_tls = os.environ.get("EDDYTOR_USE_TLS", "true").lower() != "false"
        self._skip_verify = os.environ.get("EDDYTOR_SKIP_VERIFY", "false").lower() == "true"

        self._conn = None
        self._rest: RestClient | None = None
        self._grpc: GrpcClient | None = None
        self._flight_client: flight.FlightClient | None = None

        self._connect()

    def _connect(self) -> None:
        """Establish connection to Eddytor."""
        protocol = "grpc+tls" if self._use_tls else "grpc"
        uri = f"{protocol}://{self._flight_host}:{self._flight_port}"

        db_kwargs: dict[str, str] = {
            "adbc.flight.sql.authorization_header": f"Bearer {self.api_key}"
        }

        if self._use_tls and self._skip_verify:
            db_kwargs["adbc.flight.sql.client_option.tls_skip_verify"] = "true"

        try:
            self._conn = flight_sql.connect(uri, db_kwargs=db_kwargs)
        except Exception as e:
            raise EddytorConnectionError(
                f"Failed to connect to Eddytor: {e}") from e

    @property
    def rest(self) -> RestClient:
        """
        Access the REST API client for table, storage, and domain operations.

        Returns:
            RestClient instance

        Example:
            >>> configs = client.rest.get_storage_configs()
            >>> metadata = client.rest.get_table_metadata("catalog", "schema", "table")
        """
        if self._rest is None:
            self._rest = RestClient(
                api_key=self.api_key,
                host=self._api_host,
                port=self._api_port,
                use_tls=self._use_tls,
                skip_verify=self._skip_verify,
            )
        return self._rest

    @property
    def grpc(self) -> GrpcClient:
        """
        Access the native EddytorGRPC client for structured data operations.

        Returns:
            GrpcClient instance

        Example:
            >>> result = client.grpc.stream_data("cat", "schema", "table", limit=10)
            >>> client.grpc.insert_into_table("cat", "schema", "table", arrow_table)
        """
        if self._grpc is None:
            self._grpc = GrpcClient(
                api_key=self.api_key,
                host=self._grpc_host,
                port=self._grpc_port,
                use_tls=self._use_tls,
                skip_verify=self._skip_verify,
            )
        return self._grpc

    def close(self) -> None:
        """Close the connection."""
        if self._conn:
            self._conn.close()
            self._conn = None
        if self._flight_client:
            self._flight_client.close()
            self._flight_client = None
        if self._rest:
            self._rest.close()
            self._rest = None
        if self._grpc:
            self._grpc.close()
            self._grpc = None

    def __enter__(self) -> "EddytorClient":
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.close()

    @property
    def is_connected(self) -> bool:
        """Check if client is connected."""
        return self._conn is not None

    def _ensure_connected(self) -> None:
        """Ensure connection is active."""
        if not self._conn:
            raise EddytorConnectionError(
                "Not connected. Call connect() first.")

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

        Args:
            sql: SQL query to execute

        Returns:
            pandas DataFrame with query results

        Example:
            >>> df = client.query("SELECT * FROM my_table LIMIT 10")
        """
        self._ensure_connected()

        try:
            cursor = self._conn.cursor()
            cursor.execute(sql)
            df = cursor.fetch_df()
            cursor.close()
            return df
        except Exception as e:
            raise EddytorQueryError(f"Query failed: {e}") from e

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

        Args:
            sql: SQL query to execute

        Returns:
            PyArrow Table with query results

        Example:
            >>> table = client.query_arrow("SELECT * FROM my_table")
        """
        self._ensure_connected()

        try:
            cursor = self._conn.cursor()
            cursor.execute(sql)
            table = cursor.fetch_arrow_table()
            cursor.close()
            return table
        except Exception as e:
            raise EddytorQueryError(f"Query failed: {e}") from e

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

        For DML statements (INSERT, UPDATE, DELETE), returns an empty list.
        Use execute_dml() instead if you need the affected row count.

        Args:
            sql: SQL query to execute

        Returns:
            List of tuples with query results

        Example:
            >>> rows = client.execute("SELECT id, name FROM my_table LIMIT 5")
        """
        self._ensure_connected()

        try:
            cursor = self._conn.cursor()
            cursor.execute(sql)
            if cursor.description:
                rows = cursor.fetchall()
            else:
                rows = []  # DML statement — no rows to fetch
            cursor.close()
            return rows
        except Exception as e:
            raise EddytorQueryError(f"Query failed: {e}") from e

    def execute_dml(self, sql: str) -> int:
        """
        Execute a DML statement (INSERT, UPDATE, DELETE, MERGE).

        Args:
            sql: DML SQL statement to execute

        Returns:
            Number of affected rows

        Example:
            >>> count = client.execute_dml("INSERT INTO my_table VALUES (1, 'hello')")
            >>> count = client.execute_dml("DELETE FROM my_table WHERE id = 1")
        """
        self._ensure_connected()

        try:
            cursor = self._conn.cursor()
            cursor.execute(sql)
            count = cursor.rowcount
            cursor.close()
            return count if count >= 0 else 0
        except Exception as e:
            raise EddytorQueryError(f"DML execution failed: {e}") from e

    def ingest(
        self,
        table_name: str,
        data,
        mode: str = "append",
        catalog: str | None = None,
        schema: str | None = None,
    ) -> int:
        """
        Bulk ingest Arrow data into a table via ADBC ingest (CommandStatementIngest).

        Args:
            table_name: Target table name
            data: Data to ingest (pyarrow.Table, RecordBatch, RecordBatchReader, or pandas DataFrame)
            mode: Ingest mode — "create", "append", "replace", or "create_append"
            catalog: Optional catalog name
            schema: Optional schema name

        Returns:
            Number of ingested rows

        Example:
            >>> import pyarrow as pa
            >>> table = pa.table({"id": [1, 2], "name": ["a", "b"]})
            >>> count = client.ingest("my_table", table, mode="append")
        """
        self._ensure_connected()

        try:
            cursor = self._conn.cursor()
            count = cursor.adbc_ingest(
                table_name, data, mode=mode,
                catalog_name=catalog, db_schema_name=schema,
            )
            cursor.close()
            return count
        except Exception as e:
            raise EddytorQueryError(f"Ingest failed: {e}") from e

    def list_schemas(self) -> list[str]:
        """
        List all available schemas.

        Returns:
            List of schema names
        """
        self._ensure_connected()

        try:
            cursor = self._conn.cursor()
            cursor.execute("SHOW SCHEMAS")
            rows = cursor.fetchall()
            cursor.close()
            return [row[0] for row in rows]
        except Exception as e:
            raise EddytorQueryError(f"Failed to list schemas: {e}") from e

    def list_tables(self, schema: str | None = None) -> list[str]:
        """
        List all available tables, optionally filtered by schema.

        Args:
            schema: Optional schema name to filter by

        Returns:
            List of table names
        """
        self._ensure_connected()

        try:
            cursor = self._conn.cursor()
            if schema:
                cursor.execute(f"SHOW TABLES FROM {schema}")
            else:
                cursor.execute("SHOW TABLES")
            rows = cursor.fetchall()
            cursor.close()
            return [row[0] for row in rows]
        except Exception as e:
            raise EddytorQueryError(f"Failed to list tables: {e}") from e

    def describe_table(self, table: str) -> "pd.DataFrame":
        """
        Get the schema/columns of a table.

        Args:
            table: Table name (can include schema like 'schema.table')

        Returns:
            DataFrame with column information
        """
        self._ensure_connected()

        try:
            cursor = self._conn.cursor()
            cursor.execute(f"DESCRIBE {table}")
            df = cursor.fetch_df()
            cursor.close()
            return df
        except Exception as e:
            raise EddytorQueryError(f"Failed to describe table: {e}") from e

    def table_exists(self, table: str) -> bool:
        """
        Check if a table exists.

        Args:
            table: Table name to check

        Returns:
            True if table exists, False otherwise
        """
        try:
            self.execute(f"SELECT 1 FROM {table} LIMIT 0")
            return True
        except EddytorQueryError:
            return False

    def count(self, table: str, where: str | None = None) -> int:
        """
        Count rows in a table.

        Args:
            table: Table name
            where: Optional WHERE clause (without 'WHERE' keyword)

        Returns:
            Row count

        Example:
            >>> total = client.count("my_table")
            >>> filtered = client.count("my_table", "status = 'active'")
        """
        sql = f"SELECT COUNT(*) FROM {table}"
        if where:
            sql += f" WHERE {where}"

        rows = self.execute(sql)
        return rows[0][0]

    # ------------------------------------------------------------------
    # Flight DoAction helpers (for DDL that ADBC DBAPI doesn't expose)
    # ------------------------------------------------------------------

    def _get_flight_client(self) -> flight.FlightClient:
        """Return (and lazily create) a pyarrow FlightClient."""
        if self._flight_client is None:
            scheme = "grpc+tls" if self._use_tls else "grpc"
            location = f"{scheme}://{self._flight_host}:{self._flight_port}"
            kwargs: dict[str, Any] = {}
            if self._use_tls and self._skip_verify:
                kwargs["disable_server_verification"] = True
            try:
                self._flight_client = flight.FlightClient(location, **kwargs)
            except Exception as e:
                raise EddytorConnectionError(f"Failed to create FlightClient: {e}") from e
        return self._flight_client

    def _flight_call_options(self) -> flight.FlightCallOptions:
        """Build FlightCallOptions with Bearer auth header."""
        return flight.FlightCallOptions(headers=[(b"authorization", f"Bearer {self.api_key}".encode())])

    def _do_action(self, action_type: str, body: dict[str, Any]) -> list[bytes]:
        """Send a DoAction request with a JSON-encoded body and return result buffers."""
        client = self._get_flight_client()
        action = flight.Action(action_type, json.dumps(body).encode("utf-8"))
        try:
            results = list(client.do_action(
                action, options=self._flight_call_options()))
            return [r.body.to_pybytes() for r in results]
        except flight.FlightUnavailableError as e:
            raise EddytorConnectionError(
                f"DoAction '{action_type}' failed (server unavailable): {e}") from e
        except flight.FlightError as e:
            raise EddytorQueryError(
                f"DoAction '{action_type}' failed: {e}") from e

    # ------------------------------------------------------------------
    # Arrow schema JSON conversion
    # ------------------------------------------------------------------

    @staticmethod
    def _arrow_type_to_json(arrow_type: pa.DataType) -> str | dict:
        """Convert a pyarrow DataType to the JSON representation expected by Arrow's Rust serde."""
        mapping: dict[pa.DataType, str] = {
            pa.int8(): "Int8",
            pa.int16(): "Int16",
            pa.int32(): "Int32",
            pa.int64(): "Int64",
            pa.uint8(): "UInt8",
            pa.uint16(): "UInt16",
            pa.uint32(): "UInt32",
            pa.uint64(): "UInt64",
            pa.float16(): "Float16",
            pa.float32(): "Float32",
            pa.float64(): "Float64",
            pa.string(): "Utf8",
            pa.utf8(): "Utf8",
            pa.large_string(): "LargeUtf8",
            pa.large_utf8(): "LargeUtf8",
            pa.binary(): "Binary",
            pa.large_binary(): "LargeBinary",
            pa.bool_(): "Boolean",
            pa.date32(): "Date32",
            pa.date64(): "Date64",
            pa.null(): "Null",
        }

        if arrow_type in mapping:
            return mapping[arrow_type]

        # Timestamp(unit, tz)
        if pa.types.is_timestamp(arrow_type):
            unit_map = {"s": "Second", "ms": "Millisecond",
                        "us": "Microsecond", "ns": "Nanosecond"}
            unit = unit_map[arrow_type.unit]
            tz = arrow_type.tz
            return {"Timestamp": [unit, tz]}

        # Decimal128(precision, scale)
        if pa.types.is_decimal128(arrow_type):
            return {"Decimal128": [arrow_type.precision, arrow_type.scale]}

        # Decimal256(precision, scale)
        if pa.types.is_decimal256(arrow_type):
            return {"Decimal256": [arrow_type.precision, arrow_type.scale]}

        # List(child)
        if pa.types.is_list(arrow_type):
            child_field = EddytorClient._field_to_json(arrow_type.value_field)
            return {"List": child_field}

        # LargeList(child)
        if pa.types.is_large_list(arrow_type):
            child_field = EddytorClient._field_to_json(arrow_type.value_field)
            return {"LargeList": child_field}

        # Duration(unit)
        if pa.types.is_duration(arrow_type):
            unit_map = {"s": "Second", "ms": "Millisecond",
                        "us": "Microsecond", "ns": "Nanosecond"}
            return {"Duration": unit_map[arrow_type.unit]}

        raise ValueError(f"Unsupported Arrow type: {arrow_type}")

    @staticmethod
    def _field_to_json(field: pa.Field) -> dict:
        """Convert a pyarrow Field to Arrow JSON format."""
        return {
            "name": field.name,
            "data_type": EddytorClient._arrow_type_to_json(field.type),
            "nullable": field.nullable,
            "dict_id": 0,
            "dict_is_ordered": False,
            "metadata": {k.decode() if isinstance(k, bytes) else k: v.decode() if isinstance(v, bytes) else v
                         for k, v in field.metadata.items()} if field.metadata else {},
        }

    @staticmethod
    def _schema_to_arrow_json(schema: pa.Schema) -> str:
        """Convert a pyarrow Schema to the JSON string expected by Arrow's Rust serde."""
        obj = {
            "fields": [EddytorClient._field_to_json(schema.field(i)) for i in range(len(schema))],
            "metadata": {k.decode() if isinstance(k, bytes) else k: v.decode() if isinstance(v, bytes) else v
                         for k, v in schema.metadata.items()} if schema.metadata else {},
        }
        return json.dumps(obj)

    # ------------------------------------------------------------------
    # DDL via DoAction
    # ------------------------------------------------------------------

    def create_table(
        self,
        table_name: str,
        location: str,
        schema: pa.Schema,
        constraints: list[dict[str, str]] | None = None,
        description: str | None = None,
        override_if_exists: bool = False,
    ) -> None:
        """
        Create a new Delta table via Flight SQL DoAction.

        Args:
            table_name: Name for the new table
            location: Object store URL (e.g., "s3://bucket/path")
            schema: PyArrow Schema defining the table columns
            constraints: Optional list of constraints, each {"name": "...", "expr": "..."}
            description: Optional table description
            override_if_exists: If True, overwrite existing table (default False)

        Example:
            >>> import pyarrow as pa
            >>> schema = pa.schema([
            ...     pa.field("id", pa.int64(), nullable=False),
            ...     pa.field("name", pa.string()),
            ... ])
            >>> client.create_table("my_table", "s3://bucket/my_table", schema)
        """
        body: dict[str, Any] = {
            "table_name": table_name,
            "location": location,
            "schema_json": self._schema_to_arrow_json(schema),
            "constraints": constraints or [],
            "override_if_exists": override_if_exists,
        }
        if description is not None:
            body["description"] = description

        self._do_action("eddytor.CreateTable", body)

    def add_column(
        self,
        table: str,
        columns: pa.Schema,
        catalog: str = "eddytor",
        schema: str = "eddytor",
        constraints: list[dict[str, str]] | None = None,
    ) -> None:
        """
        Add columns to an existing table via Flight SQL DoAction.

        Args:
            table: Target table name
            columns: PyArrow Schema describing the new columns to add
            catalog: Catalog name (default "eddytor")
            schema: Schema name (default "eddytor")
            constraints: Optional list of constraints, each {"name": "...", "expr": "..."}

        Example:
            >>> import pyarrow as pa
            >>> new_cols = pa.schema([pa.field("email", pa.string())])
            >>> client.add_column("my_table", new_cols)
        """
        body: dict[str, Any] = {
            "catalog": catalog,
            "schema": schema,
            "table": table,
            "schema_json": self._schema_to_arrow_json(columns),
            "constraints": constraints or [],
        }

        self._do_action("eddytor.AddColumn", body)

    def tables(self, with_discovery: bool = True) -> list[Table]:
        """
        Discover all available tables and return them as :class:`Table` handles.

        This is the recommended starting point — it calls the REST discovery
        endpoint and returns ``Table`` objects with catalog, schema, and name
        already filled in, so you never need to construct FQNs manually.

        Args:
            with_discovery: Include tables from auto-discovery (default: True).

        Returns:
            List of Table handles, one per discovered table.

        Example:
            >>> tables = client.tables()
            >>> for t in tables:
            ...     print(t.fqn)           # e.g. `eddytor`.`cfg_550e...`.`abc_customers`
            ...     print(t.count())
            >>> # Work with a specific table
            >>> customers = [t for t in tables if "customers" in t.name][0]
            >>> df = customers.query_all(limit=100)
        """
        registered = self.rest.get_registered_tables(with_discovery=with_discovery)
        return [
            Table(self, rt.catalog, rt.schema, rt.name)
            for rt in registered
            if rt.catalog and rt.schema and rt.name
        ]

    def table(self, catalog: str, schema: str, table: str) -> Table:
        """
        Get a Table handle for the given catalog.schema.name.

        If you don't know the exact catalog/schema/table values, use
        :meth:`tables` to discover them first.

        Args:
            catalog: Catalog name (always ``"eddytor"``)
            schema: Schema name (``cfg_`` + config UUID without dashes)
            table: Table name (hash prefix + name)

        Returns:
            Table handle with convenience methods

        Example:
            >>> table = client.table("eddytor", "cfg_550e8400e29b41d4a716446655440000", "abc123_products")
            >>> print(table.count())
            >>> print(table.history())
            >>> table.column("status").set_fixed_domain(["Active", "Discontinued"])
        """
        return Table(self, catalog, schema, table)

    def storage(self, config_id: str) -> Storage:
        """
        Get a Storage handle for the given config_id.

        Args:
            config_id: Storage configuration UUID

        Returns:
            Storage handle with convenience methods

        Example:
            >>> store = client.storage("cfg_xxx")
            >>> objects = store.list_objects(path="data/")
            >>> store.create_folder("data/new/")
        """
        return Storage(self, config_id)

    def get_sql_info(self) -> dict[str | int, Any]:
        """Get Flight SQL server metadata (name, version, capabilities).

        Returns a dict keyed by integer info codes or string names.
        """
        self._ensure_connected()
        try:
            return self._conn.adbc_get_info()
        except Exception as e:
            raise EddytorQueryError(f"Failed to get SQL info: {e}") from e

    def get_table_types(self) -> list[str]:
        """Get the types of tables the server supports (e.g., 'TABLE')."""
        self._ensure_connected()
        try:
            return self._conn.adbc_get_table_types()
        except Exception as e:
            raise EddytorQueryError(f"Failed to get table types: {e}") from e

    def interactive(self) -> None:
        """
        Start an interactive SQL session.

        Type SQL queries and see results. Type 'exit' to quit.
        """
        self._ensure_connected()

        print("Eddytor Interactive SQL")
        print("Type SQL queries, 'exit' to quit, 'help' for commands")
        print()

        while True:
            try:
                sql = input("SQL> ").strip()

                if not sql:
                    continue
                if sql.lower() == "exit":
                    break
                if sql.lower() == "help":
                    print("Commands:")
                    print("  SHOW SCHEMAS     - List all schemas")
                    print("  SHOW TABLES      - List all tables")
                    print("  DESCRIBE <table> - Show table schema")
                    print("  exit             - Quit interactive mode")
                    continue

                cursor = self._conn.cursor()
                cursor.execute(sql)

                if cursor.description:
                    columns = [desc[0] for desc in cursor.description]
                    print(f"Columns: {columns}")

                rows = cursor.fetchall()
                for row in rows[:50]:
                    print(row)
                if len(rows) > 50:
                    print(f"... ({len(rows)} total rows)")

                cursor.close()

            except KeyboardInterrupt:
                print("\nExiting...")
                break
            except Exception as e:
                print(f"Error: {e}")
