"""
QoScope agent with tools.
"""
import yaml
import sqlalchemy as sqa
from Gradio_UI import GradioUI
from smolagents import CodeAgent, HfApiModel, tool

from tools.diagrams import PlotTool
from tools.final_answer import FinalAnswerTool


DATABASE_PATH = 'sqlite:///resampled_daily_avg.sqlite'

db_engine = sqa.create_engine(DATABASE_PATH, echo=False)


@tool
def run_sql_query(sql_query: str) -> str:
    """
    Run a SQL query on a table in a SQLite database and return the output/result as a string.
    The output contains one row of result(s) in each line.
    The output is simple text without any Markdown styling. An agent should take this plain output
    and format appropriately when required, e.g., as a Markdown table or summarizing the results
    in words.

    The only available table in the SQLite database is `school_measurements`.
    The table's CREATE statement (DDL) is as follows:
    CREATE TABLE school_measurements (
        "index" BIGINT,
        date DATETIME,
        download_speed FLOAT,
        upload_speed FLOAT,
        latency FLOAT,
        school_id_giga TEXT,
        school_name TEXT,
        server_location TEXT,
        country TEXT,
        iso3_format TEXT
    )

    IMPORTANT: You will ONLY execute SELECT queries. You will NEVER execute any other types of
    SQL queries, e.g., INSERT, DELETE, DROP, and so on, which changes the database in any way.

    Args:
        sql_query: An appropriate, correct SQL query for a SQLite database

    Returns:
         The result of running the SQL query as a string.
    """

    # Source: https://huggingface.co/docs/smolagents/en/examples/text_to_sql

    output = ''
    with db_engine.connect() as con:
        rows = con.execute(sqa.text(sql_query))
        for row in rows:
            # Each row is a tuple
            output += '\n' + str(row)

    return output


### Main block ###

final_answer = FinalAnswerTool()
code_model = HfApiModel(
    max_tokens=2096,
    temperature=0.2,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

with open('prompts.yaml', 'r', encoding='utf-8') as stream:
    prompt_templates = yaml.safe_load(stream)

agent = CodeAgent(
    model=code_model,
    tools=[
        run_sql_query,
        PlotTool(),
        final_answer,
    ],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()