import gradio as gr
from huggingface_hub import InferenceClient
import os
from smolagents import (
    tool,
    CodeAgent,
    TransformersModel,
    GradioUI,
    MultiStepAgent,
    stream_to_gradio, HfApiModel,
)
from sqlalchemy import (
    create_engine,
    MetaData,
    Table,
    Column,
    String,
    Integer,
    Float,
    insert,
    inspect,
    text,
    select,
    Engine,
)
import spaces

from dotenv import load_dotenv

load_dotenv()
#sample questions
# What is the average each customer paid?
#  Create a sql statement and invoke your sql_engine tool


@spaces.GPU
def dummy():
    pass


@tool
def sql_engine_tool(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 = ""
    with engine.begin() as con:
        rows = con.execute(text(query))
        for row in rows:
            output += "\n" + str(row)
    return output


def init_db(engine):
    metadata_obj = MetaData()

    def insert_rows_into_table(rows, table, engine=engine):
        for row in rows:
            stmt = insert(table).values(**row)
            with engine.begin() as connection:
                connection.execute(stmt)

    table_name = "receipts"
    receipts = Table(
        table_name,
        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)

    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)

    table_name = "waiters"
    waiters = Table(
        table_name,
        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"},
    ]
    insert_rows_into_table(rows, waiters)
    return engine


if __name__ == "__main__":
    engine = create_engine("sqlite:///:localhost:")
    engine = init_db(engine)
    #Not working at the moment
    # model = TransformersModel(
    #     # model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
    #     device_map="cuda",
    #     model_id="meta-llama/Llama-3.2-3B-Instruct"
    # )
    model = HfApiModel(
          model_id="meta-llama/Llama-3.2-3B-Instruct",
        token=os.getenv("my_first_agents_hf_tokens")
    )

    agent = CodeAgent(
        tools=[sql_engine_tool],
        model=model,
        max_steps=10,
        verbosity_level=1,
    )

    def enter_message(new_message, conversation_history):
        conversation_history.append(gr.ChatMessage(role="user", content=new_message))
        # yield "", conversation_history
        for msg in stream_to_gradio(agent, new_message):
            conversation_history.append(msg)
            yield "", conversation_history


    def clear_message(chat_history: list):
        agent.memory.reset()
        return chat_history.clear(), ""

    with gr.Blocks() as b:
        gr.Markdown('''# Demo text to sql on paying customers' receipts
        a self correcting text to sql ai agent using smolagents, gradio, HF Spaces, sqlalchemy improved from a smolagents guide
        ''')
        chatbot = gr.Chatbot(type="messages", height=2000)
        message_box = gr.Textbox(lines=1, label="chat message (with default sample question)", value="What is the average each customer paid?")
        with gr.Row():
            stop_generating_button = gr.Button("stop generating")
            clear_messages_button = gr.ClearButton([message_box, chatbot])
            enter_button = gr.Button("enter")
        reply_button_click_event = enter_button.click(enter_message, [message_box, chatbot], [message_box, chatbot])
        message_submit = message_box.submit(enter_message, [message_box, chatbot], [message_box, chatbot])
        stop_generating_button.click(fn= stop_gen,cancels=[reply_button_click_event,message_submit])
        clear_messages_button.click(clear_message,outputs=[chatbot,message_box])

    b.launch()