#!/usr/bin/env python3
"""
Test utilities for Eddytor Flight SQL service.

Usage:
    # Install dependencies
    pip install eddytor-sdk

    # Run basic connectivity test
    python test_flight_sql.py --api-key edd_live_xxx

    # Run a custom query
    python test_flight_sql.py --api-key edd_live_xxx --query "SELECT * FROM my_table"

    # Test SDK Table/Column convenience API
    python test_flight_sql.py --api-key edd_live_xxx --test-sdk --catalog eddytor --db-schema cfg_xxx --table demo_products
    python test_flight_sql.py --api-key edd_live_xxx --test-sdk --catalog eddytor --db-schema cfg_xxx --table demo_products --column status
"""

import argparse
import sys

import pyarrow as pa

from eddytor_sdk import EddytorClient


def test_basic_query(client: EddytorClient):
    """Test basic SQL execution."""
    print("\n--- Basic Query Test ---")
    rows = client.execute("SELECT 1 as test_value, 'hello' as test_string")
    print(f"Result: {rows}")
    return True


def test_list_schemas(client: EddytorClient):
    """Test schema discovery."""
    print("\n--- List Schemas ---")
    try:
        schemas = client.list_schemas()
        print(f"Schemas: {schemas}")
        return True
    except Exception as e:
        print(f"Error: {e}")
        return False


def test_list_tables(client: EddytorClient):
    """List all available tables in the session."""
    print("\n--- Available Tables ---")
    try:
        tables = client.list_tables()
        if tables:
            print("Tables in current session:")
            for t in tables:
                print(f"  - {t}")
        else:
            print("No tables registered yet.")
        return True
    except Exception as e:
        print(f"Error listing tables: {e}")
        return False


def test_sdk_table(client: EddytorClient, catalog: str, schema: str, table_name: str):
    """Test the Table convenience wrapper (Flight SQL + REST)."""
    print("\n--- SDK Table Test ---")
    table = client.table(catalog, schema, table_name)
    print(f"Table handle: {table}")
    print(f"FQN:          {table.fqn}")

    # exists / count
    try:
        print(f"Exists: {table.exists()}")
    except Exception as e:
        print(f"exists() error: {e}")

    try:
        print(f"Row count: {table.count()}")
    except Exception as e:
        print(f"count() error: {e}")

    # describe
    try:
        print(f"Schema:\n{table.describe()}")
    except Exception as e:
        print(f"describe() error: {e}")

    # query_all (first 5 rows)
    try:
        df = table.query_all(limit=5)
        print(f"query_all(limit=5):\n{df}")
    except Exception as e:
        print(f"query_all() error: {e}")

    # REST: metadata
    try:
        meta = table.metadata()
        print(f"Metadata: {meta}")
    except Exception as e:
        print(f"metadata() error: {e}")

    # REST: history
    try:
        history = table.history()
        print(f"History: {history}")
    except Exception as e:
        print(f"history() error: {e}")

    print("SDK Table test complete.")
    return True


def test_sdk_column(client: EddytorClient, catalog: str, schema: str, table_name: str, column_name: str):
    """Test the Column convenience wrapper (REST domain operations)."""
    print(f"\n--- SDK Column Test ({column_name}) ---")
    col = client.table(catalog, schema, table_name).column(column_name)
    print(f"Column handle: {col}")

    # domain
    try:
        domain = col.domain()
        print(f"Domain: {domain}")
    except Exception as e:
        print(f"domain() error: {e}")

    # allowed_values
    try:
        allowed = col.allowed_values()
        print(f"Allowed values: {allowed}")
    except Exception as e:
        print(f"allowed_values() error: {e}")

    print("SDK Column test complete.")
    return True


def run_custom_query(client: EddytorClient, query: str):
    """Run a custom SQL query."""
    print(f"\n--- Custom Query ---")
    print(f"SQL: {query}")
    rows = client.execute(query)
    print(f"Rows ({len(rows)}):")
    for row in rows[:20]:
        print(f"  {row}")
    if len(rows) > 20:
        print(f"  ... and {len(rows) - 20} more")
    return rows


def interactive_mode(client: EddytorClient):
    """Interactive SQL query mode."""
    print("\n--- Interactive Mode ---")
    print("Enter SQL queries (type 'exit' to quit, 'tables' to list tables)")
    print()

    while True:
        try:
            query = input("SQL> ").strip()
            if not query:
                continue
            if query.lower() == 'exit':
                break
            if query.lower() == 'tables':
                test_list_tables(client)
                continue

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

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


def main():
    parser = argparse.ArgumentParser(description="Test Eddytor Flight SQL service")
    parser.add_argument("--api-key", required=True, help="Eddytor API key")
    parser.add_argument("--query", help="Custom SQL query to execute")
    parser.add_argument("--test-all", action="store_true", help="Run all tests")
    parser.add_argument("--interactive", "-i", action="store_true", help="Start interactive SQL mode")
    parser.add_argument("--list-tables", action="store_true", help="List available tables")
    # SDK Table/Column convenience tests
    parser.add_argument("--test-sdk", action="store_true", help="Test SDK Table/Column convenience API")
    parser.add_argument("--catalog", default="eddytor", help="Catalog name for SDK tests (default: eddytor)")
    parser.add_argument("--db-schema", help="Schema name for SDK tests (e.g. cfg_xxx)")
    parser.add_argument("--table", help="Table name for SDK tests")
    parser.add_argument("--column", help="Column name for SDK Column test")

    args = parser.parse_args()

    try:
        print("Connecting to Eddytor...")
        client = EddytorClient(api_key=args.api_key)
        print("Connected!")

        if args.test_sdk:
            if not args.db_schema or not args.table:
                print("Error: --test-sdk requires --db-schema and --table", file=sys.stderr)
                return 1
            test_sdk_table(client, args.catalog, args.db_schema, args.table)
            if args.column:
                test_sdk_column(client, args.catalog, args.db_schema, args.table, args.column)
        elif args.interactive:
            interactive_mode(client)
        elif args.list_tables:
            test_list_tables(client)
        elif args.query:
            run_custom_query(client, args.query)
        elif args.test_all:
            test_basic_query(client)
            test_list_schemas(client)
            test_list_tables(client)
        else:
            test_basic_query(client)

        client.close()
        print("\nDone!")
        return 0

    except Exception as e:
        print(f"\nError: {e}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())
