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 (in-memory) ---
engine = create_engine("sqlite:///:memory:")
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) -> str:
    """
    Executes SQL queries on the available tables: receipts and waiters.

    Args:
        query: SQL query string.
    """
    try:
        output = ""
        with engine.connect() as con:
            rows = con.execute(text(query))
            for row in rows:
                output += "\n" + str(row)
        return output or "No results."
    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"),
)

# --- Define Gradio interface ---
def ask_agent(question):
    """Ask the AI agent a question."""
    result = agent.run(question)
    return result

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()