import os
import dotenv
from smolagents import CodeAgent, ToolCallingAgent
from smolagents import OpenAIServerModel
from tools.fetch import fetch_webpage, search_web
from smolagents import VisitWebpageTool, GoogleSearchTool, SpeechToTextTool
from tools.yttranscript import get_youtube_transcript, get_youtube_title_description
import myprompts

dotenv.load_dotenv()

gemini_model = OpenAIServerModel(
    model_id="gemini-2.0-flash",
    api_key=os.environ["GEMINI_API_KEY"],
    # Google Gemini OpenAI-compatible API base URL
    api_base="https://generativelanguage.googleapis.com/v1beta/openai/",
)

vllm_model = OpenAIServerModel(
    model_id="Qwen/Qwen2.5-1.5B-Instruct",
    api_base="http://192.168.1.39:18000/v1",  
    api_key="token-abc123",  
)

openai_41nano_model = OpenAIServerModel(
    model_id="gpt-4.1-nano",
    api_base="https://api.openai.com/v1",
    api_key=os.environ["OPENAI_API_KEY"],
)


def chek_final_answer(final_answer, agent_memory)  -> bool:
    """
    Check if the final answer is correct.
    This is a placeholder function. You can implement your own logic here.
    """
    # For demonstration, we assume the answer is always correct
    return True


web_agent = CodeAgent(
    model=openai_41nano_model,
    tools=[
        search_web,
        fetch_webpage,                
    ],
    name="web_agent",
    description="Browses the web to find information",
    verbosity_level=1,
    max_steps=7,
)

audiovideo_agent = CodeAgent(
    model=openai_41nano_model,
    tools=[
        get_youtube_transcript,
        get_youtube_title_description,
        SpeechToTextTool()
    ],
    name="audiovideo_agent",
    description="Extracts information from video or audio files",
    verbosity_level=1,
    max_steps=7,
)



manager_agent = CodeAgent(
    model=openai_41nano_model,
    tools=[],
    managed_agents=[web_agent, audiovideo_agent],    
    additional_authorized_imports=["pandas", "numpy","bs4"],
    planning_interval=5,
    verbosity_level=2,
    final_answer_checks=[chek_final_answer],
    max_steps=15,
    name="manager_agent",
    description="A manager agent that coordinates the work of other agents to answer questions.",
)

class MultiAgent:
    def __init__(self):
        print("BasicAgent initialized.")

    def __call__(self, question: str) -> str:
        try:
            # log agent call in file
            with open("logs/agent_calls.log", "a") as log_file:
                log_file.write(f"Agent called with question: {question}\n")

        try:
            question = question + '\n' + myprompts.output_format
            fixed_answer = ""

            fixed_answer = manager_agent.run(question)
            
            return fixed_answer
        except Exception as e:
            error = f"An error occurred while processing the question: {e}"
            print(error)
            return error


if __name__ == "__main__":
    # Example usage
    question = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."
    agent = MultiAgent()
    answer = agent(question)
    print(f"Answer: {answer}")