"""Client for interacting with Trino server.

This module provides a client for executing queries and managing operations on Trino,
including specific support for Iceberg table operations.
"""

import json

import trino

from config import TrinoConfig


class TrinoError(Exception):
    """Base class for Trino-related errors."""

    def __init__(self, message: str):
        """Initialize with error message."""
        self.message = message
        super().__init__(self.message)


class CatalogSchemaError(TrinoError):
    """Error raised when catalog or schema information is missing."""

    def __init__(self):
        super().__init__("Both catalog and schema must be specified")


class TrinoClient:
    """A client for interacting with Trino server.

    This class provides methods to execute queries and perform administrative operations
    on a Trino server, with special support for Iceberg table operations.

    Attributes:
        config (TrinoConfig): Configuration object containing Trino connection settings.
        client (trino.dbapi.Connection): Active connection to the Trino server.
    """

    def __init__(self, config: TrinoConfig):
        """Initialize the Trino client.

        Args:
            config (TrinoConfig): Configuration object containing Trino connection settings.
        """
        self.config = config
        self.client = self._create_client()

    def _create_client(self) -> trino.dbapi.Connection:
        """Create a new Trino DB API connection.

        Returns:
            trino.dbapi.Connection: A new connection to the Trino server.
        """
        return trino.dbapi.connect(
            host=self.config.host,
            port=self.config.port,
            user=self.config.user,
            catalog=self.config.catalog,
            schema=self.config.schema,
            http_scheme=self.config.http_scheme,
            auth=self.config.auth,
            source=self.config.source,
        )

    def execute_query(self, query: str) -> str:
        """Execute a SQL query against Trino and return results as a formatted string.

        Args:
            query (str): The SQL query to execute.
            params (Optional[dict]): Dictionary of query parameters with primitive types.

        Returns:
            str: JSON-formatted string containing query results or success message.
        """
        cur: trino.dbapi.Cursor = self.client.cursor()
        cur.execute(query)
        if cur.description:
            return json.dumps(
                [dict(zip([col[0] for col in cur.description], row, strict=True)) for row in cur.fetchall()],
                default=str,
            )
        return "Query executed successfully (no results to display)"

    def get_query_history(self, limit: int) -> str:
        """Retrieve the history of executed queries.

        Args:
            limit (Optional[int]): Maximum number of queries to return. If None, returns all queries.

        Returns:
            str: JSON-formatted string containing query history.
        """
        query = "SELECT * FROM system.runtime.queries"
        if limit is not None:
            query += f" LIMIT {limit}"
        return self.execute_query(query)

    def list_catalogs(self) -> str:
        """List all available catalogs.

        Returns:
            str: Newline-separated list of catalog names.
        """
        catalogs = [row["Catalog"] for row in json.loads(self.execute_query("SHOW CATALOGS"))]
        return "\n".join(catalogs)

    def list_schemas(self, catalog: str) -> str:
        """List all schemas in a catalog.

        Args:
            catalog: The catalog name. If None, uses configured default.

        Returns:
            Newline-separated list of schema names.

        Raises:
            CatalogSchemaError: If no catalog is specified and none is configured.
        """
        catalog = catalog or self.config.catalog
        if not catalog:
            msg = "Catalog must be specified"
            raise CatalogSchemaError(msg)
        query = f"SHOW SCHEMAS FROM {catalog}"
        schemas = [row["Schema"] for row in json.loads(self.execute_query(query))]
        return "\n".join(schemas)

    def list_tables(self, catalog: str, schema: str) -> str:
        """List all tables in a schema.

        Args:
            catalog: The catalog name. If None, uses configured default.
            schema: The schema name. If None, uses configured default.

        Returns:
            Newline-separated list of table names.

        Raises:
            CatalogSchemaError: If either catalog or schema is not specified and not configured.
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            msg = "Both catalog and schema must be specified"
            raise CatalogSchemaError(msg)
        query = f"SHOW TABLES FROM {catalog}.{schema}"
        tables = [row["Table"] for row in json.loads(self.execute_query(query))]
        return "\n".join(tables)

    def describe_table(self, catalog: str, schema: str, table: str) -> str:
        """Describe the structure of a table.

        Args:
            catalog (str): The catalog name. If None, uses configured default.
            schema (str): The schema name. If None, uses configured default.
            table (str): The name of the table.

        Returns:
            str: JSON-formatted string containing table description.

        Raises:
            CatalogSchemaError: If either catalog or schema is not specified and not configured.
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            raise CatalogSchemaError
        query = f"DESCRIBE {catalog}.{schema}.{table}"
        return self.execute_query(query)

    def show_create_table(self, catalog: str, schema: str, table: str) -> str:
        """Show the CREATE TABLE statement for a table.

        Args:
            schema (str): The schema name. If None, uses configured default.
            catalog (str): The catalog name. If None, uses configured default.
            table (str): The name of the table.

        Returns:
            str: The CREATE TABLE statement for the specified table.

        Raises:
            CatalogSchemaError: If either catalog or schema is not specified and not configured.
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            raise CatalogSchemaError
        query = f"SHOW CREATE TABLE {catalog}.{schema}.{table}"
        result = json.loads(self.execute_query(query))
        return result[0]["Create Table"] if result else ""

    def show_create_view(
        self,
        catalog: str,
        schema: str,
        view: str,
    ) -> str:
        """Show the CREATE VIEW statement for a view.

        Args:
            catalog (str): The catalog name. If None, uses configured default.
            schema (str): The schema name. If None, uses configured default.
            view (str): The name of the view.

        Returns:
            str: The CREATE VIEW statement for the specified view.

        Raises:
            CatalogSchemaError: If either catalog or schema is not specified and not configured.
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            raise CatalogSchemaError
        query = f"SHOW CREATE VIEW {catalog}.{schema}.{view}"
        result = json.loads(self.execute_query(query))
        return result[0]["Create View"] if result else ""

    def show_stats(self, catalog: str, schema: str, table: str) -> str:
        """Show statistics for a table.

        Args:
            catalog (str): The catalog name. If None, uses configured default.
            schema (str): The schema name. If None, uses configured default.
            table (str): The name of the table.

        Returns:
            str: JSON-formatted string containing table statistics.

        Raises:
            CatalogSchemaError: If either catalog or schema is not specified and not configured.
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            raise CatalogSchemaError
        query = f"SHOW STATS FOR {catalog}.{schema}.{table}"
        return self.execute_query(query)

    def optimize(self, catalog: str, schema: str, table: str) -> str:
        """Optimize an Iceberg table by compacting small files.

        Args:
            catalog (str): The catalog name. If None, uses configured default.
            schema (str): The schema name. If None, uses configured default.
            table (str): The name of the table to optimize.

        Returns:
            str: Success message indicating the table was optimized.

        Raises:
            CatalogSchemaError: If either catalog or schema is not specified and not configured.
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            raise CatalogSchemaError
        query = f"ALTER TABLE {catalog}.{schema}.{table} EXECUTE optimize"
        self.execute_query(query)
        return f"Table {catalog}.{schema}.{table} optimized successfully"

    def optimize_manifests(self, table: str, catalog: str, schema: str) -> str:
        """Optimize manifest files for an Iceberg table.

        This operation reorganizes and compacts the table's manifest files for improved
        performance.

        Args:
            table (str): The name of the table.
            catalog (Optional[str]): The catalog name. If None, uses configured default.
            schema (Optional[str]): The schema name. If None, uses configured default.

        Returns:
            str: Success message indicating the manifests were optimized.

        Raises:
            CatalogSchemaError: If either catalog or schema is not specified and not configured.
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            raise CatalogSchemaError
        query = f"ALTER TABLE {catalog}.{schema}.{table} EXECUTE optimize_manifests"
        self.execute_query(query)
        return f"Manifests for table {catalog}.{schema}.{table} optimized successfully"

    def expire_snapshots(
        self,
        catalog: str,
        table: str,
        schema: str,
        retention_threshold: str = "7d",
    ) -> str:
        """Remove old snapshots from an Iceberg table.

        This operation removes snapshots older than the specified retention threshold,
        helping to manage storage and improve performance.

        Args:
            table: The name of the table.
            retention_threshold: Age threshold for snapshot removal (e.g., "7d").
            catalog: The catalog name. If None, uses configured default.
            schema: The schema name. If None, uses configured default.

        Returns:
            Success message indicating snapshots were expired.

        Raises:
            CatalogSchemaError: If either catalog or schema is not specified and not configured.
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            msg = "Both catalog and schema must be specified"
            raise CatalogSchemaError(msg)
        query = (
            f"ALTER TABLE {catalog}.{schema}.{table} "
            f"EXECUTE expire_snapshots(retention_threshold => '{retention_threshold}')"
        )
        self.execute_query(query)
        return f"Snapshots older than {retention_threshold} expired for table {catalog}.{schema}.{table}"

    def show_catalog_tree(self) -> str:
        """Show a hierarchical tree view of all catalogs, schemas, and tables.

        Returns:
            A formatted string showing the catalog > schema > table hierarchy.
        """
        tree = []
        catalogs = [row["Catalog"] for row in json.loads(self.execute_query("SHOW CATALOGS"))]
        for catalog in sorted(catalogs):
            tree.append(f"{catalog}")
            try:
                schemas = [row["Schema"] for row in json.loads(self.execute_query(f"SHOW SCHEMAS FROM {catalog}"))]
                for schema in sorted(schemas):
                    tree.append(f"{schema}")
                    try:
                        tables = [
                            row["Table"]
                            for row in json.loads(self.execute_query(f"SHOW TABLES FROM {catalog}.{schema}"))
                        ]
                        tree.extend(f" {table}" for table in sorted(tables))
                    except (trino.dbapi.TrinoQueryError, KeyError):
                        tree.append(" Unable to list tables")
            except (trino.dbapi.TrinoQueryError, KeyError):
                tree.append("Unable to list schemas")
        return "\n".join(tree) if tree else "No catalogs found"

    def show_table_properties(self, table: str, catalog: str, schema: str) -> str:
        """Show Iceberg table properties.

        Args:
            table: The name of the table
            catalog: Optional catalog name (defaults to configured catalog)
            schema: Optional schema name (defaults to configured schema)

        Returns:
            str: JSON-formatted string containing table properties
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            raise CatalogSchemaError
        query = 'SELECT * FROM "{}$properties"'
        table_identifier = f"{catalog}.{schema}.{table}"
        return self.execute_query(query.format(table_identifier))

    def show_table_history(self, table: str, catalog: str, schema: str) -> str:
        """Show Iceberg table history/changelog.

        The history contains:
        - made_current_at: TIMESTAMP(3) WITH TIME ZONE - Time when snapshot became active
        - snapshot_id: BIGINT - Identifier of the snapshot
        - parent_id: BIGINT - Identifier of the parent snapshot
        - is_current_ancestor: BOOLEAN - Whether this snapshot is an ancestor of current

        Args:
            table: The name of the table
            catalog: Optional catalog name (defaults to configured catalog)
            schema: Optional schema name (defaults to configured schema)

        Returns:
            str: JSON-formatted string containing table history
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            raise CatalogSchemaError
        table_identifier = f"{catalog}.{schema}.{table}"
        query = 'SELECT * FROM "{}$history"'
        return self.execute_query(query.format(table_identifier))

    def show_metadata_log_entries(self, table: str, catalog: str, schema: str) -> str:
        """Show Iceberg table metadata log entries.

        The metadata log contains:
        - timestamp: TIMESTAMP(3) WITH TIME ZONE - Time when metadata was created
        - file: VARCHAR - Location of the metadata file
        - latest_snapshot_id: BIGINT - ID of latest snapshot when metadata was updated
        - latest_schema_id: INTEGER - ID of latest schema when metadata was updated
        - latest_sequence_number: BIGINT - Data sequence number of metadata file

        Args:
            table: The name of the table
            catalog: Optional catalog name (defaults to configured catalog)
            schema: Optional schema name (defaults to configured schema)

        Returns:
            str: JSON-formatted string containing metadata log entries
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            raise CatalogSchemaError
        query = 'SELECT * FROM "{}$metadata_log_entries"'
        table_identifier = f"{catalog}.{schema}.{table}"
        return self.execute_query(query.format(table_identifier))

    def show_snapshots(self, table: str, catalog: str, schema: str) -> str:
        """Show Iceberg table snapshots.

        The snapshots table contains:
        - committed_at: TIMESTAMP(3) WITH TIME ZONE - Time when snapshot became active
        - snapshot_id: BIGINT - Identifier for the snapshot
        - parent_id: BIGINT - Identifier for the parent snapshot
        - operation: VARCHAR - Type of operation (append/replace/overwrite/delete)
        - manifest_list: VARCHAR - List of Avro manifest files
        - summary: map(VARCHAR, VARCHAR) - Summary of changes from previous snapshot

        Args:
            table: The name of the table
            catalog: Optional catalog name (defaults to configured catalog)
            schema: Optional schema name (defaults to configured schema)

        Returns:
            str: JSON-formatted string containing table snapshots
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            raise CatalogSchemaError
        table_identifier = f"{catalog}.{schema}.{table}$snapshots"
        query = 'SELECT * FROM "{}"'
        return self.execute_query(query.format(table_identifier))

    def show_manifests(self, table: str, catalog: str, schema: str, all_snapshots: bool = False) -> str:
        """Show Iceberg table manifests for current or all snapshots.

        The manifests table contains:
        - path: VARCHAR - Manifest file location
        - length: BIGINT - Manifest file length
        - partition_spec_id: INTEGER - ID of partition spec used
        - added_snapshot_id: BIGINT - ID of snapshot when manifest was added
        - added_data_files_count: INTEGER - Number of data files with status ADDED
        - added_rows_count: BIGINT - Total rows in ADDED files
        - existing_data_files_count: INTEGER - Number of EXISTING files
        - existing_rows_count: BIGINT - Total rows in EXISTING files
        - deleted_data_files_count: INTEGER - Number of DELETED files
        - deleted_rows_count: BIGINT - Total rows in DELETED files
        - partition_summaries: ARRAY(ROW(...)) - Partition range metadata

        Args:
            table: The name of the table
            catalog: Optional catalog name (defaults to configured catalog)
            schema: Optional schema name (defaults to configured schema)
            all_snapshots: If True, show manifests from all snapshots

        Returns:
            str: JSON-formatted string containing table manifests
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            raise CatalogSchemaError
        table_type = "all_manifests" if all_snapshots else "manifests"
        query = 'SELECT * FROM "{}${}"'
        table_identifier = f"{catalog}.{schema}.{table}"
        return self.execute_query(query.format(table_identifier, table_type))

    def show_partitions(self, table: str, catalog: str, schema: str) -> str:
        """Show Iceberg table partitions.

        The partitions table contains:
        - partition: ROW(...) - Mapping of partition column names to values
        - record_count: BIGINT - Number of records in partition
        - file_count: BIGINT - Number of files in partition
        - total_size: BIGINT - Total size of files in partition
        - data: ROW(...) - Partition range metadata with min/max values and null/nan counts

        Args:
            table: The name of the table
            catalog: Optional catalog name (defaults to configured catalog)
            schema: Optional schema name (defaults to configured schema)

        Returns:
            str: JSON-formatted string containing table partitions
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            raise CatalogSchemaError
        table_identifier = f"{catalog}.{schema}.{table}$partitions"
        query = 'SELECT * FROM "{}"'
        return self.execute_query(query.format(table_identifier))

    def show_files(self, table: str, catalog: str, schema: str) -> str:
        """Show Iceberg table data files in current snapshot.

        The files table contains:
        - content: INTEGER - Type of content (0=DATA, 1=POSITION_DELETES, 2=EQUALITY_DELETES)
        - file_path: VARCHAR - Data file location
        - file_format: VARCHAR - Format of the data file
        - record_count: BIGINT - Number of records in file
        - file_size_in_bytes: BIGINT - File size
        - column_sizes: map(INTEGER, BIGINT) - Column ID to size mapping
        - value_counts: map(INTEGER, BIGINT) - Column ID to value count mapping
        - null_value_counts: map(INTEGER, BIGINT) - Column ID to null count mapping
        - nan_value_counts: map(INTEGER, BIGINT) - Column ID to NaN count mapping
        - lower_bounds: map(INTEGER, VARCHAR) - Column ID to lower bound mapping
        - upper_bounds: map(INTEGER, VARCHAR) - Column ID to upper bound mapping
        - key_metadata: VARBINARY - Encryption key metadata
        - split_offsets: array(BIGINT) - Recommended split locations
        - equality_ids: array(INTEGER) - Field IDs for equality deletes

        Args:
            table: The name of the table
            catalog: Optional catalog name (defaults to configured catalog)
            schema: Optional schema name (defaults to configured schema)

        Returns:
            str: JSON-formatted string containing table files info
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            raise CatalogSchemaError
        table_identifier = f"{catalog}.{schema}.{table}$files"
        query = 'SELECT * FROM "{}"'
        return self.execute_query(query.format(table_identifier))

    def show_entries(self, table: str, catalog: str, schema: str, all_snapshots: bool = False) -> str:
        """Show Iceberg table manifest entries for current or all snapshots.

        The entries table contains:
        - status: INTEGER - Status of entry (0=EXISTING, 1=ADDED, 2=DELETED)
        - snapshot_id: BIGINT - ID of the snapshot
        - sequence_number: BIGINT - Data sequence number
        - file_sequence_number: BIGINT - File sequence number
        - data_file: ROW(...) - File metadata including path, format, size etc
        - readable_metrics: JSON - Human-readable file metrics

        Args:
            table: The name of the table
            catalog: Optional catalog name (defaults to configured catalog)
            schema: Optional schema name (defaults to configured schema)
            all_snapshots: If True, show entries from all snapshots

        Returns:
            str: JSON-formatted string containing manifest entries
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            raise CatalogSchemaError
        table_name = f"{catalog}.{schema}.{table}${'all_' if all_snapshots else ''}entries"
        query = 'SELECT * FROM "{}"'
        return self.execute_query(query.format(table_name))

    def show_refs(self, table: str, catalog: str, schema: str) -> str:
        """Show Iceberg table references (branches and tags).

        The refs table contains:
        - name: VARCHAR - Name of the reference
        - type: VARCHAR - Type of reference (BRANCH or TAG)
        - snapshot_id: BIGINT - ID of referenced snapshot
        - max_reference_age_in_ms: BIGINT - Max age before reference expiry
        - min_snapshots_to_keep: INTEGER - Min snapshots to keep (branches only)
        - max_snapshot_age_in_ms: BIGINT - Max snapshot age in branch

        Args:
            table: The name of the table
            catalog: Optional catalog name (defaults to configured catalog)
            schema: Optional schema name (defaults to configured schema)

        Returns:
            str: JSON-formatted string containing table references
        """
        catalog = catalog or self.config.catalog
        schema = schema or self.config.schema
        if not catalog or not schema:
            raise CatalogSchemaError
        table_identifier = f"{catalog}.{schema}.{table}$refs"
        query = 'SELECT * FROM "{}"'
        return self.execute_query(query.format(table_identifier))
