import os
from dotenv import load_dotenv
import gradio as gr
from sqlalchemy import (
    create_engine,
    MetaData,
    Table,
    Column,
    String,
    Integer,
    Float,
    insert,
    text,
)
from smolagents import tool, CodeAgent, HfApiModel

# Load Hugging Face token from environment variables
load_dotenv(override=True)
hf_token = os.getenv("HF_TOKEN")

# Initialize in-memory SQLite database
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)

# Function to insert data
def insert_rows_into_table(rows, table):
    with engine.begin() as connection:
        connection.execute(insert(table), rows)

# 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},
]
insert_rows_into_table(rows, receipts)

# SQL Execution function
@tool
def sql_engine(query: str) -> str:
    """
    Allows you to perform SQL queries on the table. Returns a string representation of the result.
    The table is named 'receipts'. Its description is as follows:
        Columns:
        - receipt_id: INTEGER
        - customer_name: VARCHAR(16)
        - price: FLOAT
        - tip: FLOAT

    Args:
        query: The query to perform. This should be correct SQL.
    """
    output = ""
    try:
        with engine.connect() as con:
            rows = con.execute(text(query))
            for row in rows:
                output += "\n" + str(row)
    except Exception as e:
        output = f"Error: {str(e)}"
    return output.strip()

# Set up the Hugging Face agent
agent = CodeAgent(
    tools=[sql_engine],
    model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct", token=hf_token),
)

# Gradio function
def query_sql(user_query):
    return sql_engine(user_query)

# Define Gradio interface
iface = gr.Interface(
    fn=query_sql,
    inputs=gr.Textbox(label="Enter your SQL Query"),
    outputs=gr.Textbox(label="Query Result"),
    title="SQL Query Executor",
    description="Enter SQL queries to interact with an in-memory SQLite database.",
    allow_flagging="never",
)

# Launch the app
if __name__ == "__main__":
    iface.launch()