import os
import gradio as gr
from smolagents import CodeAgent
from smolagents.models import InferenceClientModel
from brightdata_scraper import BrightDataScraperTool
from brightdata_search import BrightDataSearchTool
from brightdata_datasets import BrightDataDatasetTool

# Initialize tools
scraper_tool = BrightDataScraperTool()
search_tool = BrightDataSearchTool()
dataset_tool = BrightDataDatasetTool()

# Initialize the agent with a Hugging Face Inference model
# Requires HF_TOKEN in the environment for authentication.
hf_token = os.getenv("HF_TOKEN")
model = InferenceClientModel(model_id="Qwen/Qwen3-Next-80B-A3B-Thinking", token=hf_token)

agent = CodeAgent(
    tools=[scraper_tool, search_tool, dataset_tool],
    model=model,
    add_base_tools=True,
    max_steps=4,
    instructions="Answer with the first satisfactory result; do not call the same tool repeatedly once you have the needed data. Use final_answer() as soon as you can."
)


def run_agent(task: str) -> str:
    """Run the agent with the given task."""
    try:
        result = agent.run(task)
        return str(result)
    except Exception as e:
        return f"Error: {str(e)}"


# Create Gradio interface
with gr.Blocks(title="Bright Data AI Agent") as demo:
    gr.Markdown("# Bright Data AI Agent")
    gr.Markdown(
        """
    This agent can help you with web scraping, search, and quick access to Bright Data datasets.

    **Available capabilities:**
    - Search Google, Bing, or Yandex
    - Scrape any webpage (bypasses bot detection)
    - Read structured data from 40+ prebuilt datasets (e.g., amazon_product, google_maps_reviews, linkedin_company_profile)

    **Example tasks:**
    - "Search for recent AI news on Google"
    - "Scrape the content from https://example.com"
    - "Fetch google_maps_reviews for this place URL with the last 7 days"
    """
    )

    with gr.Row():
        with gr.Column():
            task_input = gr.Textbox(label="Task", placeholder="Enter your task here...", lines=3)
            submit_btn = gr.Button("Run Agent", variant="primary")

        with gr.Column():
            output = gr.Textbox(label="Result", lines=15, max_lines=30)

    submit_btn.click(fn=run_agent, inputs=[task_input], outputs=[output])

    gr.Examples(
        examples=[
            ["Search for 'latest developments in AI' on Google"],
            ["Scrape the content from https://example.com"],
            ["What are the top Python programming tutorials?"],
        ],
        inputs=[task_input],
    )


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