from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool

from Gradio_UI import GradioUI

import sqlite3
import random

# Sample data for generating employees
first_names = [
    "John",
    "Emma",
    "Michael",
    "Sarah",
    "William",
    "Olivia",
    "James",
    "Ava",
    "Robert",
    "Isabella",
]
last_names = [
    "Smith",
    "Johnson",
    "Williams",
    "Brown",
    "Jones",
    "Davis",
    "Miller",
    "Wilson",
    "Moore",
    "Taylor",
]
job_titles = [
    "Software Engineer",
    "Data Analyst",
    "Marketing Manager",
    "Sales Representative",
    "Product Designer",
]
departments = ["Engineering", "Finance", "Marketing", "Sales", "Operations"]


def create_employee_table():
    conn = sqlite3.connect("employee_database.db")  # Using in-memory database for demonstration
    cursor = conn.cursor()

    # Create employees table
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS employees (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            first_name TEXT NOT NULL,
            last_name TEXT NOT NULL,
            job_title TEXT NOT NULL,
            department_name TEXT NOT NULL,
            salary REAL NOT NULL
        )
    """
    )

    # Generate and insert sample data
    for _ in range(30):
        employee = {
            "first_name": random.choice(first_names),
            "last_name": random.choice(last_names),
            "job_title": random.choice(job_titles),
            "department_name": random.choice(departments),
            "salary": round(random.uniform(40000, 120000), 2),
        }

        cursor.execute(
            """
            INSERT INTO employees 
            (first_name, last_name, job_title, department_name, salary)
            VALUES (?, ?, ?, ?, ?)
        """,
            (
                employee["first_name"],
                employee["last_name"],
                employee["job_title"],
                employee["department_name"],
                employee["salary"],
            ),
        )

    conn.commit()

    conn.close()


# Create and populate the database
create_employee_table()


# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def get_employee_job_title(first_name:str, last_name:str)-> str: #it's import to specify the return type
    """A tool that allows to do a sqlite sql query based on first_name and last_name
    Args:
        first_name: the first argument
        last_name: the second argument
    """
    conn = sqlite3.connect('employee_database.db')
    cursor = conn.cursor()

    cursor.execute(f"SELECT job_title FROM employees WHERE first_name='{first_name}' AND last_name='{last_name}'")
    rows = cursor.fetchall()
    if len(rows) > 1:
        return "Many employees with that name"
    elif len(rows) == 0:
        return "No employee with that name"
    else:
        row = rows[0][0]
        return row

final_answer = FinalAnswerTool()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
custom_role_conversions=None,
)


# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[final_answer, get_employee_job_title], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()