from smolagents import (
    LiteLLMModel,
    OpenAIServerModel,
    ToolCallingAgent,
    CodeAgent,
    VisitWebpageTool,
    WikipediaSearchTool,
    SpeechToTextTool,
    PythonInterpreterTool,
)
import os
import yaml
import time
from typing import Optional
from tools import MySearchTool, DuckDuckGoSearchTool

google_search_tool = {"type": "google_search"}

gemini_api_key = os.environ.get("GEMINI_API_KEY", None)
open_router_api_key = os.environ.get("OPENROUTER_API_KEY", None)
gemini_model = LiteLLMModel(
    # model_id="gemini/gemini-2.0-flash",
    model_id="gemini/gemini-2.0-flash-001",
    api_key=gemini_api_key,
    tools=[google_search_tool],
)
# alternative model
model = OpenAIServerModel(
    # model_id='deepseek/deepseek-r1-0528-qwen3-8b:free',
    # model_id="thudm/glm-4-32b:free"  function calling
    model_id="google/gemini-2.0-flash-exp:free",  # function calling support
    api_key=open_router_api_key,
)

search_tool = DuckDuckGoSearchTool()
speech_tool = SpeechToTextTool()
wiki_tool = WikipediaSearchTool()
python_tool = PythonInterpreterTool()
visit_webpage_tool = VisitWebpageTool()
my_search_tool = MySearchTool()

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

prompt_templates["final_answer"] = {
    "pre_messages": "",
    "final_answer": "Here is the final answer from your managed agent '{{name}}':\n{{final_answer}}",
    "post_messages": "",
}


class SkynetMultiAgent:
    def __init__(self):
        web_agent = ToolCallingAgent(
            tools=[],
            model=gemini_model,
            name="web_search_agent",
            description="""
                This agent has access to Google Search to search for information about a specific topic.
            """,
            max_steps=2,
            prompt_templates=prompt_templates,
        )
        self.agent = CodeAgent(
            tools=[],
            model=gemini_model,
            name="manager_agent",
            description="Manages the web agents and it runs code.",
            managed_agents=[web_agent],
            prompt_templates=prompt_templates,
            max_steps=2,
        )

    def __call__(self, question: str, file_path: Optional[str] = None) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        # Adding extra content if provided
        if file_path and os.path.exists(file_path):
            eval_question = f"{question}\n\nFile provided: {file_path}"
            print(f"Processing question with file: {file_path}")
        else:
            eval_question = question

        time.sleep(2)  # 2-second delay
        fixed_answer = self.agent.run(eval_question, max_steps=5)
        print(f"Agent returning fixed answer: {fixed_answer}")
        return fixed_answer


class SkynetSingleAgent:
    def __init__(self):
        self.agent = ToolCallingAgent(
            tools=[],
            model=gemini_model,
            name="skynet_agent",
            description="Agent to answer questions and to search in the web using google search tool",
            prompt_templates=prompt_templates,
        )

    def __call__(self, question: str, file_path: Optional[str] = None) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        # Adding extra content if provided
        if file_path and os.path.exists(file_path):
            eval_question = f"{question}\n\nFile provided: {file_path}"
            print(f"Processing question with file: {file_path}")
        else:
            eval_question = question

        fixed_answer = self.agent.run(eval_question)
        print(f"Agent returning fixed answer: {fixed_answer}")
        return fixed_answer