"""
Basic usage example for the Annie SDK.

This example shows how to:
1. Set up a connector with table definitions and relationships
2. Create an agent
3. Run natural language queries (single-table and cross-table)
4. Access results
"""

import os

from annie_sdk import (
    Agent,
    ConnectorRelationship,
    ConnectorTable,
    ConnectorTableColumn,
    PostgresConnector,
)

# Set your API key (or pass it directly to Agent)
os.environ["ANNIE_API_KEY"] = "your-api-key"

# 1. Define your tables and relationships
connector = PostgresConnector(
    connection_string="postgresql://user:pass@localhost:5432/analytics",
    tables=[
        ConnectorTable(
            name="orders",
            description="Customer orders with product info",
            columns=[
                ConnectorTableColumn(name="id", type="integer"),
                ConnectorTableColumn(name="customer_name", type="string"),
                ConnectorTableColumn(name="product_id", type="integer"),
                ConnectorTableColumn(name="amount", type="decimal"),
                ConnectorTableColumn(
                    name="status",
                    type="string",
                    description="Order status: active, completed, cancelled",
                ),
                ConnectorTableColumn(name="created_at", type="timestamp"),
            ],
        ),
        ConnectorTable(
            name="products",
            description="Product catalog",
            columns=[
                ConnectorTableColumn(name="id", type="integer"),
                ConnectorTableColumn(name="name", type="string"),
                ConnectorTableColumn(name="category", type="string"),
                ConnectorTableColumn(name="price", type="decimal"),
            ],
        ),
    ],
    # Define relationships between tables for cross-table queries
    relationships=[
        ConnectorRelationship(
            primary_table="orders",
            primary_column="product_id",
            foreign_table="products",
            foreign_column="id",
            description="Each order references a product",
        ),
    ],
)

# 2. Create the agent and run queries
with Agent(connector=connector) as agent:
    # Simple query
    response = agent.run("Show me top 10 customers by total order amount")

    print(f"SQL: {response.sql}")
    print(f"Visualization: {response.visualization}")
    print(f"Success: {response.success}")

    if response.data:
        print(f"\nResults ({response.data.row_count} rows):")
        for row in response.data.rows:
            print(f"  {row}")

    # Query with AI insights
    response = agent.run("Revenue trend by month", explain=True)

    if response.data:
        print(f"\nData: {response.data.row_count} rows")

    if response.text:
        print(f"\nInsights: {response.text}")

    if response.chart:
        print(f"\nChart: {response.chart.title}")

    # Iterate over all response items
    for item in response:
        print(f"  Item type: {item.type}")

    # Serialize to JSON
    print(response.to_json(indent=2))

    # Cross-table query (uses relationships to JOIN automatically)
    response = agent.run("Show total revenue by product category")

    print(f"\nCross-table SQL: {response.sql}")
    # SQL will include: LEFT JOIN products ON orders.product_id = products.id

    if response.data:
        print(f"Results ({response.data.row_count} rows):")
        for row in response.data.rows:
            print(f"  {row}")
