import vanna as vn

# STEP 01: This is a simple example of how to use the Vanna API
api_key = vn.get_api_key('your_email')

# Set the API key and the model
vn.set_api_key(api_key)

# STEP 02:  Set the model
vn.set_model('chinook')

# STEP 03:  Connect with the database
vn.connect_to_sqlite('https://vanna.ai/Chinook.sqlite')

# STEP 04:  Ask a question
vn.ask('What are the top 10 artists by sales?')

SELECT a.name,
       sum(il.quantity) as totalsales
FROM   artist a
    INNER JOIN album al
        ON a.artistid = al.artistid
    INNER JOIN track t
        ON al.albumid = t.albumid
    INNER JOIN invoiceline il
        ON t.trackid = il.trackid
GROUP BY a.name
ORDER BY totalsales desc limit 10;

AI-generated follow-up questions:

* Who is the artist with the highest sales?
* What are the total sales for each artist?
* Which genre has the highest sales?
* Can you provide a breakdown of sales by album?
* Which artist has the lowest sales?
* What is the average sales per artist?
* Can you provide a list of the top-selling albums?
* Which genre has the least sales?
* Can you provide a breakdown of sales by country?
* What is the total sales for each genre?

# Check the models available in the account
vn.get_models()

['ecommerce-test', 'demo-tpc-h', 'tpc', 'chinook', 'thelook']

# Set the model
vn.set_model("ecommerce-test")

# Get the ddl for training the model
# Train the model with the ddl
ddl = """
CREATE TABLE if not exists stage.customers(
    customer_id           INT NOT NULL PRIMARY KEY,
    email_address         VARCHAR(50) NOT NULL,
    name                  VARCHAR(50) NOT NULL,
    business_type_id      INT NOT NULL,
    site_code             VARCHAR(10) NOT NULL,
    archived              BOOLEAN NOT NULL,
    is_key_account        BOOLEAN NOT NULL,
    date_updated          TIMESTAMP NOT NULL,
    date_created          TIMESTAMP NOT NULL,
    job_created_date  TIMESTAMP WITH TIME ZONE DEFAULT 
        CURRENT_TIMESTAMP,
    job_created_user  varchar(50) default null,
    job_updated_date  TIMESTAMP default null,
    job_updated_user  varchar(50) default null,
    CONSTRAINT fk_business_type_id FOREIGN KEY(business_type_id) REFERENCES stage.business_types (business_type_id)
);
"""
vn.train(ddl=ddl)

# Ask a question for generating the SQL
question  =  """What is the total count of new clients who registered between October 1, 2020, and 
January 1, 2022, and have made more than 10 purchases, each exceeding $20? Additionally,
 could you provide their email addresses, the number of purchases made, and the date of their
  most recent purchase?"""

print(vn.generate_sql(question=question))

SELECT COUNT(*) AS total_count,
       c.email_address,
       COUNT(o.order_id) AS num_purchases,
       MAX(o.order_date) AS most_recent_purchase_date
FROM Customers c
JOIN Orders o ON c.customer_id = o.customer_id
WHERE c.registration_date >= '2020-10-01' AND c.registration_date <= '2022-01-01'
  AND o.order_value > 20
GROUP BY c.email_address
HAVING COUNT(o.order_id) > 10;