from llama_index.core.agent.workflow import AgentWorkflow
from llama_index.tools.duckduckgo import DuckDuckGoSearchToolSpec
from llama_index.tools.wikipedia import WikipediaToolSpec
from llama_index.tools.code_interpreter import CodeInterpreterToolSpec
from llama_index.llms.google_genai import GoogleGenAI
from llama_index.core.tools import FunctionTool
from llama_index.core.workflow import Context
from llama_index.core.llms import ChatMessage
from excel_file_reader import excel_file_reader
from image_analyzer import analyze_image
from audio_analyzer import analyze_audio
from video_analyzer import analyze_video
import os

async def main(query: str, file_url: str | None) -> str: 
    google_genai_llm = GoogleGenAI(model="gemini-2.0-flash", api_key=os.getenv("GOOGLE_API_KEY"))

    # Create tool instances
    tools = []
    tools.append(FunctionTool.from_defaults(DuckDuckGoSearchToolSpec().duckduckgo_full_search, name="duckduckgo_search", description="Searches the web using DuckDuckGo."))
    tools.append(FunctionTool.from_defaults(WikipediaToolSpec().search_data, name="wikipedia_search", description="Searches Wikipedia for information."))
    tools.append(FunctionTool.from_defaults(CodeInterpreterToolSpec().code_interpreter, name="code_interpreter", description="Executes Python code and returns the output."))

    image_analyzer_tool = FunctionTool.from_defaults(analyze_image, name="analyze_image", description="Analyzes an image provided by URL using Google Gemini API. Creates summary of the image content.")
    tools.append(image_analyzer_tool)

    audio_analyzer_tool = FunctionTool.from_defaults(analyze_audio, name="analyze_audio", description="Transcribes audio files to text using Google Gemini Flash model. Supports various audio formats including MP3, WAV, M4A, etc.")
    tools.append(audio_analyzer_tool)

    video_analyzer_tool = FunctionTool.from_defaults(analyze_video, name="analyze_video", description="Analyzes a YouTube video using Google Gemini API. Creates summary of the video content.")
    tools.append(video_analyzer_tool)

    excel_file_reader_tool = FunctionTool.from_defaults(excel_file_reader, name="excel_file_reader", description="Loads Excel files and analyzes the content of the file.")
    tools.append(excel_file_reader_tool)
    
    agent = AgentWorkflow.from_tools_or_functions(  
        tools_or_functions=tools,
        llm=google_genai_llm,
        system_prompt="""
    You are a general AI assistant. I will ask you a question. 
    Report your thoughts, and finish your answer with the following 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.
""" 
    )

    ctx = Context(agent)

    prompt = query if file_url is None else f"{query}. Use the URL {file_url} as the Url path of attached files mentioned in the question. Attached file are required to answer the question. To run Python code use Code Intrepreter tool."
    message = ChatMessage(
        role="user",
        content=prompt
    )

    response = await agent.run(message, ctx=ctx)

    return response