import time
import requests

# --- Configuration ---
DATABRICKS_INSTANCE = "dbc-da15e658-afcd.cloud.databricks.com"
API_TOKEN = "dapie17b5e771da1b7b481063e027e570e13"
SPACE_ID = "01f1580f04271dc0aa6ad4b41298f0e4"
USER_PROMPT = "What were our top 5 selling products last month by total revenue?"

# Base setups
BASE_URL = f"https://{DATABRICKS_INSTANCE}/api/2.0/genie/spaces/{SPACE_ID}"
HEADERS = {
    "Authorization": f"Bearer {API_TOKEN}",
    "Content-Type": "application/json"
}

def fetch_table_data(statement_id):
    """Fetches the actual rows and data for the query Genie executed."""
    print(f"Fetching table results for Statement: {statement_id}...")
    sql_url = f"https://{DATABRICKS_INSTANCE}/api/2.0/sql/statements/{statement_id}"
    
    # Poll the SQL execution warehouse if it's still finalizing the chunk
    while True:
        res = requests.get(sql_url, headers=HEADERS)
        res.raise_for_status()
        sql_data = res.json()
        status = sql_data.get("status", {}).get("state")
        
        if status == "SUCCEEDED":
            manifest = sql_data.get("manifest", {})
            columns = [col["name"] for col in manifest.get("schema", {}).get("columns", [])]
            data_array = sql_data.get("result", {}).get("data_array", [])
            
            # Print data cleanly in a list format
            print(f"\nColumns: {columns}")
            print("-" * 50)
            for row in data_array:
                print(row)
            break
        elif status in ["FAILED", "CANCELED"]:
            print(f"Failed to extract structured warehouse data: {status}")
            break
        time.sleep(2)

def ask_genie():
    # Step 1: Start a new conversation thread inside the Space
    print("Initializing Genie conversation thread...")
    start_url = f"{BASE_URL}/start-conversation"
    start_payload = {"content": USER_PROMPT}
    
    response = requests.post(start_url, headers=HEADERS, json=start_payload)
    response.raise_for_status()
    
    convo_data = response.json()
    # Handle both payload response flavors
    conversation_id = convo_data.get("conversation_id")
    message_id = convo_data.get("message_id") or convo_data.get("id")
    
    print(f"Conversation Started! Convo ID: {conversation_id}")
    
    # Step 2 & 3: Poll the execution status until complete
    status_url = f"{BASE_URL}/conversations/{conversation_id}/messages/{message_id}"
    
    print("Waiting for Genie to process metadata and execute SQL queries...")
    while True:
        status_response = requests.get(status_url, headers=HEADERS)
        status_response.raise_for_status()
        
        message_status = status_response.json()
        current_state = message_status.get("status")
        
        print(f"Current Status: {current_state}")
        
        if current_state == "COMPLETED":
            print("\n--- Genie Summary Generated ---")
            
            # 1. Print Text Response Summary
            attachments = message_status.get("attachments", [])
            for attachment in attachments:
                if "text" in attachment:
                    print(attachment["text"].get("content"))
            
            # 2. Extract Data Table Results using Statement ID
            query_metadata = message_status.get("query_result_metadata")
            if query_metadata and "statement_id" in query_metadata:
                fetch_table_data(query_metadata["statement_id"])
            else:
                # Alternate location Databricks drops query details inside the attachments nested structures
                for attachment in attachments:
                    if "query" in attachment and "query_result_metadata" in attachment["query"]:
                        stmt_id = attachment["query"]["query_result_metadata"].get("statement_id")
                        if stmt_id:
                            fetch_table_data(stmt_id)
            break
            
        elif current_state in ["FAILED", "CANCELLED", "QUERY_RESULT_EXPIRED"]:
            print(f"\nExecution terminated with status: {current_state}")
            if "error" in message_status:
                print(f"Error Details: {message_status['error']}")
            break
            
        time.sleep(3)

if __name__ == "__main__":
    ask_genie()