from __future__ import annotations

import gradio as gr

from tool import BrightDataSearchTool

ENGINE_CHOICES = ("google", "bing", "yandex")
tool = BrightDataSearchTool()


def run(query: str, engine: str) -> str:
    return tool(query=query, engine=engine)


def create_demo() -> gr.Blocks:
    with gr.Blocks() as demo:
        gr.Markdown("### Bright Data search")
        gr.Markdown("### [Signup here](https://brightdata.com/?utm_source=tech-partner&utm_medium=link&utm_campaign=huggingface&hs_signup=1)")
        query = gr.Textbox(label="Query")
        engine = gr.Dropdown(choices=list(ENGINE_CHOICES), label="Engine", value=ENGINE_CHOICES[0])
        run_btn = gr.Button("Run")
        output = gr.Textbox(label="Output", lines=12)
        run_btn.click(run, inputs=[query, engine], outputs=output)
    return demo


if __name__ == "__main__":
    create_demo().launch()