import json
import yaml
from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Load prompt templates from prompts.yaml
try:
    with open("prompts.yaml", 'r') as stream:
        prompt_templates = yaml.safe_load(stream)
except Exception as e:
    print(f"⚠️ Error loading prompt templates: {e}")
    prompt_templates = {}

# Research Tool (Limits articles to avoid hitting API limits)
@tool
def limited_research_tool(query: str, max_articles: int = 3) -> str:
    """Fetches only the first few research articles to avoid search limit errors.

    Args:
        query: The research topic to search online.
        max_articles: Number of articles to retrieve and summarize (default: 3).

    Returns:
        A summarized version of the key insights from the articles.
    """
    print(f"🔍 Researching: {query} (Max {max_articles} articles)")

    research_tool = DuckDuckGoSearchTool()  # Ensure tool is initialized

    try:
        # Get search results (limit articles)
        search_results = research_tool.run(query)[:max_articles]

        # Summarize articles (Placeholder summarizer)
        summaries = [f"Summary of {article}" for article in search_results]

        return "\n\n".join(summaries) if summaries else "No useful research found."

    except Exception as e:
        print(f"❌ Research failed: {e}")
        return "Research unavailable at the moment."

# Writer Tool (Formats research into an article draft)
@tool
def writer_tool(research_data: str) -> str:
    """
    Creates a structured article draft based on research data.

    Args:
        research_data: A string containing summarized research insights, typically 
                       generated by the limited_research_tool function. Each summary 
                       is separated by two newline characters and represents a key 
                       point or source.

    Returns:
        A draft article in string format, structured using the provided research data.
    """
    return f"📄 Draft article:\n{research_data}"


# Editor Tool (Refines and finalizes the article)
@tool
def editor_tool(draft: str) -> str:
    """
    Refines and enhances the draft article.

    Args:
        draft: A string representing the initial draft of an article, usually generated 
               by the writer_tool. This draft includes summarized research content 
               and serves as the basis for polishing the final article.

    Returns:
        A string containing the improved and refined version of the article.
    """
    return f"✍️ Edited article:\n{draft}"


# Final Answer Tool (Ensures structured output)
final_answer = FinalAnswerTool()

# Define LLM Model
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
)

# Researcher Agent (Uses Limited Research Tool)
researcher = CodeAgent(
    model=model,
    tools=[limited_research_tool],
    max_steps=6,
    verbosity_level=2,
    name="Researcher",
    description=prompt_templates.get("managed_agent", {}).get("researcher", {}).get("task", "Research topic details."),
    prompt_templates=prompt_templates,
)

# Writer Agent
writer = CodeAgent(
    model=model,
    tools=[writer_tool],
    max_steps=6,
    verbosity_level=2,
    name="Writer",
    description=prompt_templates.get("managed_agent", {}).get("writer", {}).get("task", "Write an article draft."),
    prompt_templates=prompt_templates,
)

# Editor Agent (Includes Final Answer Tool)
editor = CodeAgent(
    model=model,
    tools=[editor_tool, final_answer],
    max_steps=6,
    verbosity_level=2,
    name="Editor",
    description=prompt_templates.get("managed_agent", {}).get("editor", {}).get("task", "Edit and finalize an article."),
    prompt_templates=prompt_templates,
)

# Orchestrator Agent (Manages Research, Writing, and Editing)
orchestrator_agent = CodeAgent(
    model=model,
    tools=[limited_research_tool, writer_tool, editor_tool, final_answer],
    max_steps=6,
    verbosity_level=2,
    name="Orchestrator",
    description="An agent that coordinates research, writing, and editing.",
    prompt_templates=prompt_templates,
)

# Launch Gradio UI
print("🚀 Launching Gradio UI...")
GradioUI(orchestrator_agent).launch()