from smolagents import OpenAIServerModel, CodeAgent, InferenceClientModel, DuckDuckGoSearchTool, VisitWebpageTool
import markdownify
import os

import tools
import prompts

MANAGER_MODEL_GPT = "gpt-4.5-preview"
FINAL_ANSWER_MODEL_GEMINI = "gemini-2.5-pro-preview-03-25" 
AGENT_MODEL_GTP = "gpt-4.1-mini"

MANAGER_MODEL = "deepseek-ai/DeepSeek-R1"
# FINAL_ANSWER_MODEL = "gpt-4o" # OpenAIServerModel
FINAL_ANSWER_MODEL = "deepseek-ai/DeepSeek-R1" # OpenAIServerModel
AGENT_MODEL = "Qwen/Qwen2.5-Coder-32B-Instruct"
WEB_SEARCH_MODEL        = "Qwen/Qwen2.5-Coder-32B-Instruct"
IMAGE_ANALYSIS_MODEL    = "HuggingFaceM4/idefics2-8b"
AUDIO_ANALYSIS_MODEL    = "Qwen/Qwen2-Audio-7B-Instruct"
VIDEO_ANALYSIS_MODEL    = "llava-hf/LLaVA-NeXT-Video-7B-hf"
YOUTUBE_ANALYSIS_MODEL  = "llava-hf/LLaVA-NeXT-Video-7B-hf"
DOCUMENT_ANALYSIS_MODEL = "Qwen/Qwen2.5-Coder-32B-Instruct"
ARITHMETIC_MODEL        = "Qwen/Qwen2.5-Coder-32B-Instruct"
CODE_GENERATION_MODEL   = "Qwen/Qwen2.5-Coder-32B-Instruct"
CODE_EXECUTION_MODEL    = "Qwen/Qwen2.5-Coder-32B-Instruct"

# Agents

def create_custom_web_search_agent(message):
    return CodeAgent(
        name="custom_web_search_agent",
        description=prompts.get_web_search_prompt(message),
        model=InferenceClientModel(WEB_SEARCH_MODEL),
        max_steps=3,
        tools=[tools.simple_web_search_tool, tools.visit_web_page_tool],
    )

def create_simple_web_search_agent(message):
    return CodeAgent(
        name="simple_web_search_agent",
        description=prompts.get_web_search_prompt(message),
        model=InferenceClientModel(WEB_SEARCH_MODEL),
        max_steps=3,
        tools=[tools.simple_web_search_tool, tools.visit_web_page_tool],
    )

def create_image_analysis_agent(message):
    return CodeAgent(
        name="image_analysis_agent",
        description=prompts.get_image_analysis_prompt(message),
        model=InferenceClientModel(IMAGE_ANALYSIS_MODEL),
        tools=[tools.image_analysis_tool],
        max_steps=3,
    )

def create_audio_analysis_agent(message):
    return CodeAgent(
        name="audio_analysis_agent",
        description=prompts.get_audio_analysis_prompt(message),
        model=InferenceClientModel(AUDIO_ANALYSIS_MODEL),
        tools=[tools.audio_analysis_tool],
        max_steps=3,
    )

def create_video_analysis_agent(message):
    return CodeAgent(
        name="video_analysis_agent",
        description=prompts.get_video_analysis_prompt(message),
        model=InferenceClientModel(VIDEO_ANALYSIS_MODEL),
        tools=[tools.video_analysis_tool],
        max_steps=3,
    )

def create_youtube_analysis_agent(message):
    return CodeAgent(
        name="youtube_analysis_agent",
        description=prompts.get_youtube_analysis_prompt(message),
        model=InferenceClientModel(YOUTUBE_ANALYSIS_MODEL),
        tools=[tools.youtube_analysis_tool],
        max_steps=3,
    )

def create_document_analysis_agent(message):
    return CodeAgent(
        name="document_analysis_agent",
        description=prompts.get_document_analysis_prompt(message),
        model=InferenceClientModel(DOCUMENT_ANALYSIS_MODEL),
        tools=[tools.document_analysis_tool],
        max_steps=3,
    )

def create_arithmetic_agent(message):
    return CodeAgent(
        name="arithmetic_agent",
        description=prompts.get_arithmetic_prompt(message),
        model=InferenceClientModel(ARITHMETIC_MODEL),
        tools=[
            tools.add,
            tools.subtract,
            tools.multiply,
            tools.divide,
            tools.modulus,
        ],
        max_steps=3,
    )

def create_code_generation_agent(message):
    return CodeAgent(
        name="code_generation_agent",
        description=prompts.get_code_generation_prompt(message),
        model=InferenceClientModel(CODE_GENERATION_MODEL),
        tools=[tools.code_generation_tool],
        max_steps=3,
    )

def create_code_execution_agent(message):
    return CodeAgent(
        name="code_execution_agent",
        description=prompts.get_code_execution_prompt(message),
        model=InferenceClientModel(CODE_EXECUTION_MODEL),
        tools=[tools.code_execution_tool],
        max_steps=3,
    )

def create_manager_agent(message):
    simple_web_search_agent = create_simple_web_search_agent(message)
    image_analysis_agent = create_image_analysis_agent(message)
    audio_analysis_agent = create_audio_analysis_agent(message)
    video_analysis_agent = create_video_analysis_agent(message)
    youtube_analysis_agent = create_youtube_analysis_agent(message)
    document_analysis_agent = create_document_analysis_agent(message)
    arithmetic_agent = create_arithmetic_agent(message)
    code_generation_agent = create_code_generation_agent(message)
    code_execution_agent = create_code_execution_agent(message)

    return CodeAgent(
        name="manager_agent",
        model=InferenceClientModel(MANAGER_MODEL, provider="together", max_tokens=8096),
        description=prompts.get_manager_prompt(message),
        tools=[],
        planning_interval=4,
        verbosity_level=2,
        managed_agents=[
            simple_web_search_agent,
            image_analysis_agent,
            audio_analysis_agent,
            video_analysis_agent,
            youtube_analysis_agent,
            document_analysis_agent,
            arithmetic_agent,
            code_generation_agent,
            code_execution_agent,
        ],
        max_steps=10,
        additional_authorized_imports=[
            "requests",
            "zipfile",
            "os",
            "pandas",
            "numpy",
            "sympy",
            "json",
            "bs4",
            "pubchempy",
            "xml",
            "yahoo_finance",
            "Bio",
            "sklearn",
            "scipy",
            "pydub",
            "io",
            "PIL",
            "chess",
            "PyPDF2",
            "pptx",
            "torch",
            "datetime",
            "csv",
            "fractions",
        ],
    )

def create_final_answer_agent(message):
    return CodeAgent(
        name="final_answer_agent",
        description="Given a question and an initial answer, return the final refined answer following strict formatting rules.",
        # model=OpenAIServerModel(FINAL_ANSWER_MODEL),
        model=InferenceClientModel(FINAL_ANSWER_MODEL),
        tools=[],
    )