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

from Gradio_UI import GradioUI
import gradio as gr

import re
from markdownify import markdownify
from requests.exceptions import RequestException

with open('styles.css', 'r', encoding='utf-8') as file:
    styles = file.read()

with open('scripts.js', 'r', encoding='utf-8') as file:
    scripts = """
        async () => {
        function waitForElementById(id, callback) {
            const element = document.getElementById(id);
            if (element) {
                callback(element);
                return;
            }
            const observer = new MutationObserver((mutationsList, observer) => {
                for (let mutation of mutationsList) {
                    if (mutation.type === 'childList') {
                const element = document.getElementById(id);
                        if (element) {
                            callback(element);
                            observer.disconnect();
                            return;
                        }
                    }
                }
            });
                observer.observe(document.body, { childList: true, subtree: true });
            }
            waitForElementById('root', () => {
                %s
            });
            
            }
        """ % (file.read())

@tool
def json_string_to_poster(json_string:str)-> str: 
    """A tool that generate a poster from stringified json object 
    Args:
        json_string: stringified json data
    """

    return f"""
        <script id="cheatsheets_json">{repair_json(json_string)}</script>
    """

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
model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',
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, DuckDuckGoSearchTool(), json_string_to_poster],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

def get_and_visualize_data(query):
    response = agent.run(query)
    return response
    

# 4. Создаем Gradio интерфейс
with gr.Blocks(css=styles) as demo:
    with gr.Row():
        prompt_input = gr.Textbox(label="prompt", placeholder="Example: create poster with npm commands")
    
    poster_display = gr.HTML("""
        <div id="root"></div>
    """);

    json_html = gr.HTML();

    prompt_input.submit(
        get_and_visualize_data,
        prompt_input,
        json_html,
    )

    demo.load(None, None, None, js=scripts)

demo.launch()