import os
from smolagents import (
    CodeAgent,
    OpenAIServerModel,
    ToolCallingAgent,
    ToolCollection,
    DuckDuckGoSearchTool, 
    InferenceClientModel,
    GoogleSearchTool,
    VisitWebpageTool, 
    tool,
    LiteLLMModel,
)

from opentelemetry.sdk.trace import TracerProvider

from openinference.instrumentation.smolagents import SmolagentsInstrumentor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from mcp import StdioServerParameters
from tools.transcribe import load_images

trace_provider = TracerProvider()
trace_provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter()))

SmolagentsInstrumentor().instrument(tracer_provider=trace_provider)


from  tools.mathTools import multiply, add, subtract, divide, modulus
from tools.searchTools import wiki_search, mini_web_search, arvix_search
from tools.fileTool import download_file, text_file_tool, download_and_read_excel_file
from tools.transcribe import parse_youtube_video
from tools.transcribe_audio import transcribe_mp3_with_whisper
def build_agents():
    #"deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
    #deepseek-ai/DeepSeek-R1-0528
    model = OpenAIServerModel(model_id="gpt-4o")
    #model = InferenceClientModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct") #Qwen/Qwen2.5-Coder-32B-Instruct")
#     server_parameters = StdioServerParameters(
#     command="uvx",
#     args=["--quiet", "pubmedmcp@0.1.3"],
#     env={"UV_PYTHON": "3.12", **os.environ},
# )

    # with ToolCollection.from_mcp(server_parameters, trust_remote_code=True) as tool_collection:
    agent = CodeAgent(
        tools=[
            #*tool_collection.tools,
            GoogleSearchTool(provider="serper"),
            VisitWebpageTool(),
            multiply, 
            add, 
            subtract, 
            divide, 
            modulus,
            wiki_search, 
            arvix_search,
            download_file,
            text_file_tool,
            download_and_read_excel_file,
            parse_youtube_video,
            transcribe_mp3_with_whisper], 
        model=model, 
        additional_authorized_imports=["time","pandas","json","numpy","markdownify","requests","re","openpyxl","beautifulsoup4"],
        step_callbacks=[load_images],
        planning_interval=3,
        max_steps=10,
        add_base_tools=True)

  
    prompt = f"""
    FACTS: Remember that if you have to do a grouping about FOOD, then ICE CREAM is also a FOOD.

    PLEASE FOLLOW THE INSTRUCTIONS FOR ANSWERING CAREFULLY:
        Your answer should follow the template: FINAL ANSWER: [YOUR FINAL ANSWER].
        YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
        If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. 
        If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. 
        If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."""
    

    agent.prompt_templates["system_prompt"] = agent.prompt_templates["system_prompt"] + prompt


    return agent