#!/usr/bin/env python3
"""
Eddytor SDK — Flight SQL examples.

Prerequisites:
    pip install eddytor-sdk

Usage:
    # Run all examples
    python examples_flight_sql.py --api-key edd_live_xxx

    # Run a specific example
    python examples_flight_sql.py --api-key edd_live_xxx --example query

    # List available examples
    python examples_flight_sql.py --list
"""

from __future__ import annotations

import argparse
import sys

# ---------------------------------------------------------------------------
# SDK import — the only dependency beyond Python stdlib
# ---------------------------------------------------------------------------
sys.path.insert(0, ".")  # allow running from tools/python/
from eddytor_sdk import EddytorClient, EddytorQueryError


# ===========================================================================
# Example helpers
# ===========================================================================

def heading(title: str) -> None:
    print(f"\n{'='*60}")
    print(f"  {title}")
    print(f"{'='*60}\n")


def success(msg: str) -> None:
    print(f"  [OK] {msg}")


def fail(msg: str) -> None:
    print(f"  [FAIL] {msg}")


# ===========================================================================
# 1. Basic queries
# ===========================================================================

def example_query(client: EddytorClient) -> None:
    """Run SELECT queries and get results as DataFrames or raw tuples."""
    heading("Example 1: Basic Queries")

    # --- 1a. Simple expression ---
    print("1a. SELECT literal values:")
    rows = client.execute("SELECT 1 AS id, 'hello' AS greeting")
    print(f"    Result: {rows}")
    success("execute() returns list[tuple]")

    # --- 1b. Query as Arrow Table ---
    print("\n1b. Same query as Arrow Table:")
    table = client.query_arrow("SELECT 1 AS id, 'hello' AS greeting")
    print(f"    Schema: {table.schema}")
    print(f"    Rows:   {table.num_rows}")
    print(f"    Data:   {table.to_pydict()}")
    success("query_arrow() returns pyarrow.Table")

    # --- 1c. Query as pandas DataFrame ---
    try:
        print("\n1c. Same query as pandas DataFrame:")
        df = client.query("SELECT 42 AS answer, 'world' AS name")
        print(f"    {df.to_string(index=False)}")
        success("query() returns pandas.DataFrame")
    except ImportError:
        print("    (skipped — pandas not installed)")


# ===========================================================================
# 2. List schemas and tables
# ===========================================================================

def example_list_tables(client: EddytorClient) -> None:
    """Discover what schemas and tables exist in the session."""
    heading("Example 2: List Schemas & Tables")

    # --- 2a. Schemas ---
    print("2a. Available schemas:")
    try:
        schemas = client.list_schemas()
        for s in schemas:
            print(f"    - {s}")
        success(f"Found {len(schemas)} schema(s)")
    except EddytorQueryError as e:
        fail(f"list_schemas: {e}")

    # --- 2b. Tables ---
    print("\n2b. Available tables:")
    try:
        tables = client.list_tables()
        if tables:
            for t in tables:
                print(f"    - {t}")
        else:
            print("    (none — tables appear after first query via storage config)")
        success(f"Found {len(tables)} table(s)")
    except EddytorQueryError as e:
        fail(f"list_tables: {e}")


# ===========================================================================
# 3. Describe a table
# ===========================================================================

def example_describe(client: EddytorClient) -> None:
    """Show a table's schema using DESCRIBE."""
    heading("Example 3: Describe Table")

    # Replace with an actual table name from your session:
    table_name = "your_catalog.your_schema.your_table"

    print(f"Describing '{table_name}'...")
    print("(Replace table_name in the example with a real table from your session)\n")

    try:
        df = client.describe_table(table_name)
        print(df.to_string(index=False))
        success("describe_table() works")
    except EddytorQueryError as e:
        print(f"    Expected error (table doesn't exist): {e}")
        print("    Tip: Run example_list_tables first to find available tables.")


# ===========================================================================
# 4. Count rows
# ===========================================================================

def example_count(client: EddytorClient) -> None:
    """Count rows, optionally with a filter."""
    heading("Example 4: Count Rows")

    table_name = "your_catalog.your_schema.your_table"

    print(f"Counting rows in '{table_name}'...")
    print("(Replace table_name with a real table from your session)\n")

    try:
        total = client.count(table_name)
        print(f"    Total rows: {total}")
        success("count() without filter")
    except EddytorQueryError as e:
        print(f"    Expected error (table doesn't exist): {e}")

    # With a WHERE filter:
    # filtered = client.count(table_name, where='"status" = \'active\'')
    # print(f"    Active rows: {filtered}")


# ===========================================================================
# 5. DML — INSERT, UPDATE, DELETE
# ===========================================================================

def example_dml(client: EddytorClient) -> None:
    """Execute DML statements and get affected row counts."""
    heading("Example 5: DML (INSERT / UPDATE / DELETE)")

    table_name = "your_catalog.your_schema.your_table"

    print("NOTE: Replace table_name and column names with your actual table.\n")
    print("These examples show the API — they will error without a real table.\n")

    # --- 5a. INSERT ---
    print("5a. INSERT a single row:")
    print('    client.execute_dml("INSERT INTO my_table VALUES (100, \'New Row\')")')
    try:
        count = client.execute_dml(
            f"INSERT INTO {table_name} VALUES (100, 'New Row')"
        )
        success(f"Inserted {count} row(s)")
    except EddytorQueryError as e:
        print(f"    Error (expected without real table): {e}")

    # --- 5b. UPDATE ---
    print("\n5b. UPDATE rows:")
    print('    client.execute_dml("UPDATE my_table SET name = \'Updated\' WHERE id = 100")')
    try:
        count = client.execute_dml(
            f'UPDATE {table_name} SET "Name" = \'Updated\' WHERE "PK" = 100'
        )
        success(f"Updated {count} row(s)")
    except EddytorQueryError as e:
        print(f"    Error (expected without real table): {e}")

    # --- 5c. DELETE ---
    print("\n5c. DELETE rows:")
    print('    client.execute_dml("DELETE FROM my_table WHERE id = 100")')
    try:
        count = client.execute_dml(
            f'DELETE FROM {table_name} WHERE "PK" = 100'
        )
        success(f"Deleted {count} row(s)")
    except EddytorQueryError as e:
        print(f"    Error (expected without real table): {e}")

    # --- 5d. execute() also works for DML (returns empty list) ---
    print("\n5d. execute() also handles DML gracefully (returns []):")
    print('    rows = client.execute("DELETE FROM my_table WHERE id = 999")')
    print("    # Returns [] for DML, no crash")


# ===========================================================================
# 6. Bulk ingest via ADBC
# ===========================================================================

def example_ingest(client: EddytorClient) -> None:
    """Bulk-upload data using ADBC ingest (CommandStatementIngest)."""
    heading("Example 6: Bulk Ingest (ADBC)")

    import pyarrow as pa

    table_name = "your_table"

    # Build a PyArrow Table
    data = pa.table({
        "PK":   pa.array([10, 20, 30], type=pa.int64()),
        "Name": pa.array(["Alice", "Bob", "Carol"], type=pa.string()),
    })

    print(f"Ingesting {data.num_rows} rows into '{table_name}':")
    print(f"  Schema: {data.schema}")
    print(f"  Data:   {data.to_pydict()}")

    try:
        count = client.ingest(
            table_name,
            data,
            mode="append",        # "create", "append", "replace", "create_append"
            catalog="datafusion",  # optional
            schema="public",       # optional
        )
        success(f"Ingested {count} row(s)")
    except EddytorQueryError as e:
        print(f"\n    Error (expected without real table): {e}")

    # --- Ingest from pandas ---
    print("\n  You can also pass a pandas DataFrame directly:")
    print('    df = pd.DataFrame({"PK": [1, 2], "Name": ["X", "Y"]})')
    print('    client.ingest("my_table", df, mode="append")')


# ===========================================================================
# 7. Check table existence
# ===========================================================================

def example_table_exists(client: EddytorClient) -> None:
    """Check if a table exists before querying it."""
    heading("Example 7: Table Existence Check")

    print("Checking if 'nonexistent_table' exists...")
    exists = client.table_exists("nonexistent_table")
    print(f"    Exists: {exists}")
    success(f"table_exists() returned {exists}")


# ===========================================================================
# 8. Context manager usage
# ===========================================================================

def example_context_manager(api_key: str) -> None:
    """Use the SDK as a context manager for automatic cleanup."""
    heading("Example 8: Context Manager")

    print("Using 'with' statement for automatic cleanup:\n")
    print("    with EddytorClient(api_key=...) as client:")
    print('        df = client.query("SELECT 1")')
    print("    # connection is closed automatically\n")

    with EddytorClient(api_key=api_key) as client:
        rows = client.execute("SELECT 'context_manager_works' AS status")
        print(f"    Result: {rows}")
        success("Context manager works — connection auto-closed on exit")


# ===========================================================================
# 9. Error handling
# ===========================================================================

def example_error_handling(client: EddytorClient) -> None:
    """Show how to handle errors from the SDK."""
    heading("Example 9: Error Handling")

    from eddytor_sdk import EddytorQueryError

    # --- 9a. Bad SQL ---
    print("9a. Handling bad SQL:")
    try:
        client.execute("THIS IS NOT VALID SQL")
    except EddytorQueryError as e:
        print(f"    Caught EddytorQueryError: {e}")
        success("Bad SQL is caught cleanly")

    # --- 9b. Nonexistent table ---
    print("\n9b. Querying a nonexistent table:")
    try:
        client.query("SELECT * FROM does_not_exist_xyz")
    except EddytorQueryError as e:
        print(f"    Caught EddytorQueryError: {e}")
        success("Missing table is caught cleanly")

    # --- 9c. DML on nonexistent table ---
    print("\n9c. DML on nonexistent table:")
    try:
        client.execute_dml("INSERT INTO does_not_exist_xyz VALUES (1)")
    except EddytorQueryError as e:
        print(f"    Caught EddytorQueryError: {e}")
        success("DML errors are caught cleanly")


# ===========================================================================
# 10. Full workflow: query -> transform -> write back
# ===========================================================================

def example_full_workflow(client: EddytorClient) -> None:
    """End-to-end workflow: read data, transform, write back."""
    heading("Example 10: Full Read-Transform-Write Workflow")

    print("Typical workflow with a real table:\n")
    print("""
    import pyarrow as pa
    import pyarrow.compute as pc

    # 1. Read data
    table = client.query_arrow("SELECT * FROM sales WHERE region = 'EU'")
    print(f"Read {table.num_rows} rows")

    # 2. Transform with PyArrow compute
    total = pc.sum(table.column("amount"))
    print(f"Total EU sales: {total}")

    # 3. Create a summary table
    summary = pa.table({
        "region": ["EU"],
        "total_amount": [total.as_py()],
        "row_count": [table.num_rows],
    })

    # 4. Write the summary back via ingest
    client.ingest("sales_summary", summary, mode="append")

    # 5. Or insert via SQL DML
    client.execute_dml(
        f"INSERT INTO sales_log VALUES ('EU', {total.as_py()}, {table.num_rows})"
    )
    """)
    success("See code above for the pattern")


# ===========================================================================
# Main runner
# ===========================================================================

EXAMPLES = {
    "query":          ("Basic queries (SELECT)",                example_query),
    "list":           ("List schemas & tables",                 example_list_tables),
    "describe":       ("Describe a table's schema",             example_describe),
    "count":          ("Count rows",                            example_count),
    "dml":            ("DML (INSERT/UPDATE/DELETE)",             example_dml),
    "ingest":         ("Bulk ingest (ADBC)",                    example_ingest),
    "exists":         ("Table existence check",                 example_table_exists),
    "context":        ("Context manager usage",                 None),  # special handling
    "errors":         ("Error handling",                        example_error_handling),
    "workflow":       ("Full read-transform-write workflow",    example_full_workflow),
}


def main():
    parser = argparse.ArgumentParser(
        description="Eddytor SDK — Flight SQL examples",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Run all examples
  python examples_flight_sql.py --api-key edd_live_xxx

  # Run just the query examples
  python examples_flight_sql.py --api-key edd_live_xxx --example query
""",
    )
    parser.add_argument("--api-key", required=True, help="Eddytor API key")
    parser.add_argument(
        "--example",
        choices=list(EXAMPLES.keys()),
        help="Run a specific example",
    )
    parser.add_argument(
        "--list-examples", action="store_true", help="List available examples"
    )

    args = parser.parse_args()

    if args.list_examples:
        print("Available examples:\n")
        for name, (desc, _) in EXAMPLES.items():
            print(f"  {name:<12}  {desc}")
        print(f"\nRun one:  python {sys.argv[0]} --api-key KEY --example query")
        print(f"Run all:  python {sys.argv[0]} --api-key KEY")
        return 0

    print("Connecting to Eddytor...")
    client = EddytorClient(api_key=args.api_key)
    print("Connected!\n")

    passed = 0
    failed = 0

    def run(name: str):
        nonlocal passed, failed
        desc, fn = EXAMPLES[name]
        try:
            if name == "context":
                example_context_manager(args.api_key)
            else:
                fn(client)
            passed += 1
        except Exception as e:
            fail(f"{desc}: {e}")
            import traceback
            traceback.print_exc()
            failed += 1

    if args.example:
        run(args.example)
    else:
        for name in EXAMPLES:
            run(name)

    client.close()

    print(f"\n{'='*60}")
    print(f"  Results: {passed} passed, {failed} failed")
    print(f"{'='*60}")
    return 1 if failed else 0


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