"""
Convexity SDK - Query the DuckLake Warehouse

This sample demonstrates how to:
- Retrieve read-only DuckLake connection info for a project
- Attach the DuckLake warehouse using DuckDB (local Python driver)
- Run SQL queries against your project's datasets

Prerequisites:
- A valid API key with VIEWER (or higher) permission
- The "Convexity SDK samples" org and "griffin it" project must exist.
  Run the previous samples first:
    1. create_org_and_project.py
    2. create_databricks_connections.py
    3. create_datasets.py
- ``duckdb`` Python package installed:
    pip install duckdb   # or:  uv pip install duckdb
"""

import sys

import duckdb

from convexity_sdk import ConvexityClient

ORG_SLUG = "convexity-sdk-samples"
PROJECT_SLUG = "griffin-it"

# ── Initialise SDK client ───────────────────────────────────────────────
client = ConvexityClient()

# ── Step 1: Verify organization exists ──────────────────────────────────
org = client.organizations.get_by_slug(ORG_SLUG)
if org is None or not hasattr(org, "id"):
    print(
        f'ERROR: Organization "{ORG_SLUG}" not found.\nPlease run create_org_and_project.py first.',
        file=sys.stderr,
    )
    raise SystemExit(1)
print(f"Found organization: {org.name} (id={org.id})")

# ── Step 2: Verify project exists ──────────────────────────────────────
project = client.projects.get_by_slug(org.id, PROJECT_SLUG)
if project is None or not hasattr(project, "id"):
    print(
        f'ERROR: Project "{PROJECT_SLUG}" not found in org "{ORG_SLUG}".\nPlease run create_org_and_project.py first.',
        file=sys.stderr,
    )
    raise SystemExit(1)
print(f"Found project: {project.name} (id={project.id})")

# ── Step 3: Get datalake connection info ────────────────────────────────
print("\nFetching datalake connection info …")

dl_resp = client.datalake.get_info(project_id=project.id)
if dl_resp is None or not dl_resp.available or dl_resp.connection is None:
    msg = getattr(dl_resp, "message", None) or "Datalake not available for this project."
    print(f"ERROR: {msg}", file=sys.stderr)
    raise SystemExit(1)

conn_info = dl_resp.connection
print(f"  Backend type : {conn_info.backend_type}")

if conn_info.backend_type.value != "ducklake":
    print(
        f"ERROR: This sample only supports the 'ducklake' backend, but got '{conn_info.backend_type.value}'.",
        file=sys.stderr,
    )
    raise SystemExit(1)

# ── Step 4: Attach the DuckLake warehouse with DuckDB ──────────────────
print("\nAttaching DuckLake warehouse via DuckDB …")

db = duckdb.connect()

# Install and load extensions (httpfs must be ready before ducklake accesses S3)
db.execute("INSTALL httpfs; LOAD httpfs;")
db.execute("INSTALL ducklake; LOAD ducklake;")

# Create an S3 secret so DuckDB can access the datalake storage
# Strip protocol prefix — DuckDB expects bare hostname; USE_SSL handles HTTPS
endpoint = conn_info.s_3_endpoint or ""
endpoint = endpoint.replace("https://", "").replace("http://", "")

secret_sql = f"""CREATE SECRET convexity_s3 (
    TYPE S3,
    KEY_ID '{conn_info.s_3_access_key_id}',
    SECRET '{conn_info.s_3_secret_access_key}',
    ENDPOINT '{endpoint}',
    USE_SSL true,
    URL_STYLE 'path',
    REGION '{conn_info.s_3_region or "auto"}'
);"""
db.execute(secret_sql)

ducklake_name = conn_info.ducklake_name or "datalake"

# Attach using the catalog path (S3) in read-only mode
attach_sql = f"ATTACH 'ducklake:{conn_info.catalog_path}' AS {ducklake_name} (DATA_PATH '{conn_info.data_path}', READ_ONLY);"
print(f"  Running SQL: {attach_sql}")
db.execute(attach_sql)
print(f"  Attached DuckLake as '{ducklake_name}'")

# Switch to the attached database
db.execute(f"USE {ducklake_name};")

# ── Step 5: Run example queries ─────────────────────────────────────────
print("\n── Example queries ──────────────────────────────────────────")

# Discover tables directly from the attached DuckLake database
tables_result = db.execute("SHOW TABLES;").fetchall()
table_names = [row[0] for row in tables_result]

print(f"\nTables in '{ducklake_name}':")
for name in table_names:
    print(f"  - {name}")

if not table_names:
    print("\n  No tables available yet. Run create_datasets.py to populate data.")
else:
    # Query the first available table (preview first 5 rows)
    first_table = table_names[0]
    query = f"SELECT * FROM {ducklake_name}.{first_table} LIMIT 5;"

    print(f"\nPreview of '{first_table}' (first 5 rows):")
    print(f"  SQL: {query}\n")

    result = db.execute(query)
    columns = [desc[0] for desc in result.description]
    rows = result.fetchall()

    # Print column headers
    header = " | ".join(f"{c:>20}" for c in columns)
    print(f"  {header}")
    print(f"  {'-' * len(header)}")

    # Print rows
    for row in rows:
        formatted = " | ".join(f"{str(v):>20}" for v in row)
        print(f"  {formatted}")

    print(f"\n  ({len(rows)} rows returned)")

# ── Cleanup ─────────────────────────────────────────────────────────────
db.close()
client.close()
print("\nDone!")
