"""
AIStor Tables Basic Example with PyIceberg

This example demonstrates fundamental operations with AIStor Tables using PyIceberg:
- Connecting to an AIStor Tables catalog
- Loading an existing table
- Appending data to a table
- Running various types of queries

Prerequisites:
- AIStor server running with Tables feature enabled
- A warehouse, namespace, and table already created
- PyIceberg and PyArrow installed: pip install pyiceberg pyarrow pandas

Configuration:
- Update CATALOG_URL, WAREHOUSE, NAMESPACE, and TABLE_NAME below
- Set your access credentials in ACCESS_KEY and SECRET_KEY
"""

import os

import pandas as pd
import pyarrow as pa
import pyiceberg.catalog.rest
from pyiceberg.catalog import Catalog
from pyiceberg.table import Table

# =============================================================================
# Configuration - Update these values for your environment
# =============================================================================

CATALOG_URL = "http://localhost:9000/_iceberg"
HOST = "http://localhost:9000"
WAREHOUSE = "mywarehouse"
NAMESPACE = "mynamespace"
TABLE_NAME = "mytable"
ACCESS_KEY = "minioadmin"
SECRET_KEY = "minioadmin"
AWS_REGION = "local"


def setup_environment(access_key: str, secret_key: str, region: str) -> None:
    """Configure AWS environment variables for authentication."""
    os.environ["AWS_ACCESS_KEY_ID"] = access_key
    os.environ["AWS_SECRET_ACCESS_KEY"] = secret_key
    os.environ["AWS_REGION"] = region


def create_catalog(
    catalog_url: str,
    warehouse: str,
    access_key: str,
    secret_key: str,
    host: str,
    region: str,
) -> Catalog:
    """
    Create a PyIceberg REST catalog connected to AIStor Tables.

    Args:
        catalog_url: URL of the AIStor Iceberg REST catalog endpoint
        warehouse: Name of the warehouse (table bucket)
        access_key: Access key for authentication
        secret_key: Secret key for authentication
        host: S3 endpoint host URL
        region: AWS region for signing requests

    Returns:
        Configured RestCatalog instance
    """
    config = {
        # Enable SigV4 authentication for the REST catalog
        "rest.sigv4-enabled": "true",
        "rest.signing-name": "s3tables",
        "rest.signing-region": region,
        # S3 configuration for data access
        "s3.access-key-id": access_key,
        "s3.secret-access-key": secret_key,
        "s3.endpoint": host,
        "s3.path-style-access": "true",
    }

    catalog = pyiceberg.catalog.rest.RestCatalog(
        name="aistor_catalog",
        uri=catalog_url,
        warehouse=warehouse,
        **config,
    )

    print(f"Connected to catalog at {catalog_url}")
    return catalog


def load_table(catalog: Catalog, namespace: str, table_name: str) -> Table:
    """
    Load an existing table from the catalog.

    Args:
        catalog: PyIceberg catalog instance
        namespace: Namespace containing the table
        table_name: Name of the table to load

    Returns:
        Loaded Table instance
    """
    table_identifier = (namespace, table_name)
    table = catalog.load_table(table_identifier)
    print(f"Loaded table: {namespace}.{table_name}")
    return table


def append_data(table: Table, data: dict) -> None:
    """
    Append data to an Iceberg table.

    The data dictionary keys must match the table schema field names.
    PyArrow infers the schema from the table.

    Args:
        table: Target Iceberg table
        data: Dictionary with column names as keys and lists as values
    """

    # Define the table schema
    # id is required, other fields optional
    schema = pa.schema([
        pa.field("id", pa.int64(), nullable=False),
        pa.field("name", pa.string()),
        pa.field("category", pa.string()),
        pa.field("price", pa.float64()),
        pa.field("in_stock", pa.bool_()),
    ])

    # Create a PyArrow table using this schema
    arrow_table = pa.Table.from_pydict(data, schema=schema)
    
    # Append to the Iceberg table
    table.append(arrow_table)

    row_count = len(next(iter(data.values())))
    print(f"Appended {row_count} rows to table")


# =============================================================================
# Query Examples
# =============================================================================


def query_all_rows(table: Table) -> pd.DataFrame:
    """
    Scan all rows from the table.

    Returns:
        DataFrame containing all rows
    """
    print("\n--- Query: All Rows ---")
    df = table.scan().to_pandas()
    print(f"Found {len(df)} total rows")
    print(df)
    return df


def query_with_filter(table: Table, filter_expression: str) -> pd.DataFrame:
    """
    Query rows matching a filter condition.

    Args:
        table: Table to query
        filter_expression: SQL-like filter (e.g., "price > 100")

    Returns:
        DataFrame containing matching rows
    """
    print(f"\n--- Query: Filter '{filter_expression}' ---")
    df = table.scan(row_filter=filter_expression).to_pandas()
    print(f"Found {len(df)} matching rows")
    print(df)
    return df


def query_selected_columns(
    table: Table, columns: tuple, filter_expression: str = None
) -> pd.DataFrame:
    """
    Query specific columns, optionally with a filter.

    Args:
        table: Table to query
        columns: Tuple of column names to retrieve
        filter_expression: Optional filter condition

    Returns:
        DataFrame containing selected columns
    """
    print(f"\n--- Query: Columns {columns} ---")

    scan = table.scan(selected_fields=columns)
    if filter_expression:
        scan = table.scan(selected_fields=columns, row_filter=filter_expression)
        print(f"    Filter: {filter_expression}")

    df = scan.to_pandas()
    print(f"Found {len(df)} rows")
    print(df)
    return df


def query_with_limit(table: Table, limit: int) -> pd.DataFrame:
    """
    Query a limited number of rows.

    Args:
        table: Table to query
        limit: Maximum number of rows to return

    Returns:
        DataFrame containing up to 'limit' rows
    """
    print(f"\n--- Query: Limit {limit} rows ---")
    df = table.scan(limit=limit).to_pandas()
    print(f"Retrieved {len(df)} rows")
    print(df)
    return df


def display_table_info(table: Table) -> None:
    """Display basic table metadata and schema."""
    print("\n--- Table Information ---")
    print(f"Name: {table.name()}")
    print(f"Location: {table.location()}")
    print("\nSchema:")
    for field in table.schema().fields:
        nullable = "nullable" if field.optional else "required"
        print(f"  {field.name}: {field.field_type} ({nullable})")


def display_snapshot_info(table: Table) -> None:
    """Display current snapshot information."""
    print("\n--- Snapshot Information ---")
    snapshot = table.current_snapshot()
    if snapshot:
        print(f"Snapshot ID: {snapshot.snapshot_id}")
        print(f"Timestamp: {snapshot.timestamp_ms}")
        if snapshot.summary:
            print(f"Operation: {snapshot.summary.operation}")
    else:
        print("No snapshots available")


# =============================================================================
# Main Example
# =============================================================================


def main():
    """
    Main function demonstrating AIStor Tables operations.

    This example assumes:
    - The warehouse already exists
    - The namespace already exists
    - The table already exists with a compatible schema

    To create these resources, use the AIStor admin tools or a setup script.
    """
    print("=" * 60)
    print("AIStor Tables Basic Example")
    print("=" * 60)

    # Step 1: Configure environment and connect to catalog
    setup_environment(ACCESS_KEY, SECRET_KEY, AWS_REGION)
    catalog = create_catalog(
        CATALOG_URL, WAREHOUSE, ACCESS_KEY, SECRET_KEY, HOST, AWS_REGION
    )

    # Step 2: Load an existing table
    table = load_table(catalog, NAMESPACE, TABLE_NAME)

    # Step 3: Display table information
    display_table_info(table)

    # Step 4: Append sample data
    # Note: Column names and types must match the existing table schema
    sample_data = {
        "id": [101, 102, 103],
        "name": ["Widget A", "Widget B", "Widget C"],
        "category": ["Electronics", "Electronics", "Hardware"],
        "price": [29.99, 49.99, 19.99],
        "in_stock": [True, True, False],
    }

    print("\n" + "=" * 60)
    print("Appending Data")
    print("=" * 60)
    append_data(table, sample_data)

    # Refresh table to see the new snapshot
    table = load_table(catalog, NAMESPACE, TABLE_NAME)

    # Step 5: Run various queries
    print("\n" + "=" * 60)
    print("Query Examples")
    print("=" * 60)

    # Query all rows
    query_all_rows(table)

    # Query with filter: price greater than 30
    query_with_filter(table, "price > 30")

    # Query with filter: in stock items only
    query_with_filter(table, "in_stock = true")

    # Query with filter: specific category
    query_with_filter(table, "category = 'Electronics'")

    # Query selected columns only
    query_selected_columns(table, ("name", "price"))

    # Query selected columns with filter
    query_selected_columns(
        table, ("name", "price", "in_stock"), "category = 'Electronics'"
    )

    # Query with limit
    query_with_limit(table, 5)

    # Step 6: Display snapshot info after changes
    print("\n" + "=" * 60)
    print("Snapshot After Append")
    print("=" * 60)
    display_snapshot_info(table)

    print("\n" + "=" * 60)
    print("Example Complete")
    print("=" * 60)


if __name__ == "__main__":
    main()