import os
import gradio as gr
from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, Float, insert, text, inspect
from smolagents import tool, CodeAgent, InferenceClientModel

# --- Setup SQLite database (persistent in Space) ---
engine = create_engine("sqlite:///data.db")
metadata_obj = MetaData()

# --- Create receipts table ---
receipts = Table(
    "receipts",
    metadata_obj,
    Column("receipt_id", Integer, primary_key=True),
    Column("customer_name", String(16), primary_key=True),
    Column("price", Float),
    Column("tip", Float),
)
metadata_obj.create_all(engine)

# Insert sample data
rows = [
    {"receipt_id": 1, "customer_name": "Alan Payne", "price": 12.06, "tip": 1.20},
    {"receipt_id": 2, "customer_name": "Alex Mason", "price": 23.86, "tip": 0.24},
    {"receipt_id": 3, "customer_name": "Woodrow Wilson", "price": 53.43, "tip": 5.43},
    {"receipt_id": 4, "customer_name": "Margaret James", "price": 21.11, "tip": 1.00},
]
for row in rows:
    stmt = insert(receipts).values(**row)
    with engine.begin() as conn:
        conn.execute(stmt)

# --- Create waiters table ---
waiters = Table(
    "waiters",
    metadata_obj,
    Column("receipt_id", Integer, primary_key=True),
    Column("waiter_name", String(16), primary_key=True),
)
metadata_obj.create_all(engine)

rows = [
    {"receipt_id": 1, "waiter_name": "Corey Johnson"},
    {"receipt_id": 2, "waiter_name": "Michael Watts"},
    {"receipt_id": 3, "waiter_name": "Michael Watts"},
    {"receipt_id": 4, "waiter_name": "Margaret James"},
]
for row in rows:
    stmt = insert(waiters).values(**row)
    with engine.begin() as conn:
        conn.execute(stmt)

# --- Define SQL tool for the agent ---
@tool
def sql_engine(query: str) -> list:
    """
    Executes SQL queries on the available tables: receipts and waiters.

    Args:
        query: SQL query string.

    Returns:
        List of tuples with query results.
    """
    try:
        print(f"🧩 Executing query: {query}")  # Debug log
        with engine.connect() as con:
            rows = con.execute(text(query))
            results = [tuple(row) for row in rows]  # Keep results as tuples
        return results or []
    except Exception as e:
        return [f"⚠️ SQL Error: {str(e)}"]

# Dynamically describe tables
updated_description = "Allows SQL queries on the following tables:\n"
inspector = inspect(engine)
for table in ["receipts", "waiters"]:
    columns_info = [(col["name"], col["type"]) for col in inspector.get_columns(table)]
    table_description = f"\n\nTable '{table}':\nColumns:\n" + "\n".join(
        [f" - {name}: {col_type}" for name, col_type in columns_info]
    )
    updated_description += table_description

sql_engine.description = updated_description

# --- Create the agent ---
agent = CodeAgent(
    tools=[sql_engine],
    model=InferenceClientModel(
        "meta-llama/Meta-Llama-3-8B-Instruct",
        api_key=os.environ.get("HF_TOKEN")  # Use secret from Space
    ),
)

# --- Define Gradio interface ---
def ask_agent(question: str) -> str:
    """Ask the AI agent a question and return its answer."""
    try:
        result = agent.run(question)
        # Convert tuples to readable string if result is a list of tuples
        if isinstance(result, list):
            result_str = "\n".join(str(r) for r in result)
            return result_str or "No results."
        return str(result)
    except Exception as e:
        return f"⚠️ Error: {str(e)}"

demo = gr.Interface(
    fn=ask_agent,
    inputs=gr.Textbox(label="Ask a question (e.g., Who got the biggest tip?)"),
    outputs=gr.Textbox(label="Agent response"),
    title="🧠 Text-to-SQL Agent",
    description="Ask natural-language questions about a restaurant receipts database. Powered by smolagents + Meta-Llama-3.",
)

demo.launch()