import streamlit as st
from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, ManagedAgent, VisitWebpageTool, tool


from sqlalchemy import (
    create_engine,
    MetaData,
    Table,
    Column,
    String,
    Integer,
    Float,
    insert,
    inspect,
    text,
)

engine = create_engine("sqlite:///:memory:")
metadata_obj = MetaData()

# create city SQL table
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},
]
for row in rows:
    stmt = insert(receipts).values(**row)
    with engine.begin() as connection:
        cursor = connection.execute(stmt)

inspector = inspect(engine)
columns_info = [(col["name"], col["type"]) for col in inspector.get_columns("receipts")]

table_description = "Columns:\n" + "\n".join([f"  - {name}: {col_type}" for name, col_type in columns_info])
print(table_description)

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


sql_agent = CodeAgent(
    tools=[sql_engine],
    model=HfApiModel("meta-llama/Meta-Llama-3.1-8B-Instruct"),
)

managed_sql_agent = ManagedAgent(
    agent=sql_agent,
    name="sql search",
    description="Runs SQL queries. Give it your query as an argument. Also, this agent should provide the actual query it ran.",
)

manager_agent = CodeAgent(
    tools=[],
    model=HfApiModel(model_id = "Qwen/Qwen2.5-Coder-32B-Instruct"),
    managed_agents=[managed_web_agent],
    additional_authorized_imports=['pyparsing', 'matplotlib', 'datetime', 'statistics', 'bs4', 'request', 'unicodedata', 'queue', 'time', 'collections', 're', 'math', 'stat', 'random', 
'itertools'],
)

# Function to log agent actions
def log_agent_action(prompt, result, agent_name):
    st.write(f"### Agent Activity ({agent_name}):")
    st.write("**Prompt Sent to Agent:**")
    st.code(prompt, language="text")
    st.write("**Agent Output:**")
    st.code(result, language="text")

# Streamlit app title
st.title("AI SQL Assistant Agent researching your query and summarizing it")

# App description
st.write("Generate SQL queries using human speech powered by SmolAgents.")

# Input blog topic or prompt

prompt = st.text_area("Enter the your prompt:", placeholder="E.g., Can you give me the name of the client who got the most expensive receipt?")

# Button to generate blog content
if st.button("Generate Summary"):
    if prompt:
        with st.spinner("Generating content..."):
            try:
                # Run the blog agent with the given prompt
                result = manager_agent.run(prompt)
                # Display the generated blog content
                st.subheader("Generated Summary:")
                st.write(result)

                # Log backend activity
                log_agent_action(prompt, result, "SQL Agent")
            except Exception as e:
                st.error(f"An error occurred: {e}")
    else:
        st.warning("Please enter a blog topic or prompt to proceed.")

# Footer
st.markdown("---")
st.caption("Powered by SmolAgents, and Streamlit")