import os
from smolagents import CodeAgent, OpenAIServerModel, PythonInterpreterTool
from tools.fetch import fetch_webpage, search_web
from tools.yttranscript import get_youtube_transcript, get_youtube_title_description
from tools.stt import get_text_transcript_from_audio_file
from tools.image import analyze_image
from common.mylogger import mylog
from myprompts import output_format
import dotenv

dotenv.load_dotenv()

# --- Models ---
openai_model = OpenAIServerModel(
    model_id="gpt-3.5-turbo",
    api_key=os.environ["OPENAI_API_KEY"],
    api_base="https://api.openai.com/v1",
)

# --- Tool Agents ---
web_agent = CodeAgent(
    model=openai_model,
    tools=[search_web, fetch_webpage],
    name="web_agent",
    description="Searches the web and fetches content from webpages.",
    max_steps=6,
    verbosity_level=1,
)

audiovideo_agent = CodeAgent(
    model=openai_model,
    tools=[
        get_youtube_transcript,
        get_youtube_title_description,
        get_text_transcript_from_audio_file,
        analyze_image
    ],
    name="audiovideo_agent",
    description="Processes audio, video, and image content.",
    max_steps=6,
    verbosity_level=1,
)

# --- Answer Checker ---
def check_final_answer(answer: str, _) -> bool:
    return isinstance(answer, str) and len(answer.strip()) <= 200

# --- Manager Agent ---
manager_agent = CodeAgent(
    model=openai_model,
    tools=[PythonInterpreterTool()],
    managed_agents=[web_agent, audiovideo_agent],
    name="manager_agent",
    description="Manages sub-agents to answer complex queries.",
    final_answer_checks=[check_final_answer],
    verbosity_level=2,
    max_steps=12,
)

# --- Entry Class ---
class MultiAgent:
    def __init__(self):
        print("✅ MultiAgent initialized")

    def __call__(self, question: str) -> str:
        mylog("MultiAgent", question)
        try:
            system_instruction = (
                "You are the lead agent in a multi-agent AI system. Use the sub-agents wisely."
                " Answer using reasoning, planning, and tools where needed."
                " Stay concise and always format the answer exactly as required."
            )

            prompt = f"""
{system_instruction}

THE QUESTION:
{question.strip()}

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


if __name__ == "__main__":
    example_question = "What was the actual enrollment of the Malko competition in 2023?"
    agent = MultiAgent()
    print(agent(example_question))