from smolagents import CodeAgent, HfApiModel, load_tool, tool
import datetime
import pytz
import yaml
from transformers import pipeline
from duckduckgo_search import DDGS
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Real summarization model
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")

# Real DuckDuckGo search + summarization tool
@tool
def search_and_summarize(query: str, max_sentences: int) -> str:
    """A tool that performs a web search and returns a summary.
    Args:
        query: The search query string.
        max_sentences: Maximum number of sentences in the summary.
    """
    try:
        with DDGS() as ddgs:
            results = ddgs.text(query, max_results=5)
            search_texts = [r["body"] for r in results if "body" in r]
            combined_text = "\n".join(search_texts)

        if not combined_text.strip():
            return "No content found to summarize."

        summary = summarizer(
            combined_text,
            max_length=60 * max_sentences,
            min_length=20 * max_sentences,
            do_sample=False,
        )[0]["summary_text"]
        return summary
    except Exception as e:
        return f"Error during search or summarization: {str(e)}"

# Timezone tool
@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"

# Final answer tool
final_answer = FinalAnswerTool()

# Model configuration
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
    custom_role_conversions=None,
)

# Load image tool (optional)
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

# Load system prompts
with open("prompts.yaml", "r") as stream:
    prompt_templates = yaml.safe_load(stream)

# Agent setup
agent = CodeAgent(
    model=model,
    tools=[search_and_summarize, get_current_time_in_timezone, final_answer],
    max_steps=6,
    verbosity_level=2,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates,
)

# Register tools for use in python_interpreter
globals()["search_and_summarize"] = search_and_summarize
globals()["get_current_time_in_timezone"] = get_current_time_in_timezone

# Launch the UI
GradioUI(agent).launch()