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
from datetime import datetime


# Connect to SQLite database. It will be created if it doesn't exist
conn = sqlite3.connect('employee_database.db')
cursor = conn.cursor()

# Create employees table with basic columns
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
)
''')

# Commit the transaction
conn.commit()

# Sample data for employees
employees_data = [
    ('John', 'Smith', 'Software Engineer'),
    ('Sarah', 'Johnson', 'Marketing Manager'),
    ('Michael', 'Brown', 'Data Scientist'),
    ('Emma', 'Wilson', 'Product Designer'),
    ('William', 'Davis', 'DevOps Engineer'),
    ('Olivia', 'Miller', 'UX Researcher'),
    ('James', 'Wilson', 'Backend Developer'),
    ('Jessica', 'Moore', 'Project Manager'),
    ('Robert', 'Taylor', 'Frontend Developer'),
    ('Lisa', 'Anderson', 'Quality Assurance'),
    ('Richard', 'Thomas', 'Technical Writer'),
    ('Amy', 'Jackson', 'Business Analyst'),
    ('Charles', 'White', 'Database Administrator'),
    ('Elizabeth', 'Harris', 'Customer Support'),
    ('Thomas', 'Martin', 'Network Engineer'),
    ('Jennifer', 'Thompson', 'HR Specialist'),
    ('Donald', 'Walker', 'Security Specialist'),
    ('Margaret', 'Young', 'Content Manager'),
    ('Ronald', 'Allen', 'Sales Representative'),
    ('Susan', 'Scott', 'Financial Analyst')
]

# Insert multiple records efficiently using executemany
cursor.executemany('''
INSERT INTO employees (first_name, last_name, job_title)
VALUES (?, ?, ?)
''', employees_data)

# Commit the transaction
conn.commit()

# 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="https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud/",
custom_role_conversions=None,
)

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