# agent.py — full GAIA-ready agent with working WikipediaQueryRun + tools

import os
import asyncio

from llama_index.llms.openai import OpenAI
from llama_index.core.agent.react.base import ReActAgent
from llama_index.core.tools import FunctionTool

from langchain_community.tools.wikipedia.tool import WikipediaQueryRun
from langchain_community.utilities.wikipedia import WikipediaAPIWrapper
from langchain_experimental.tools.python.tool import PythonREPLTool
from langchain_community.document_loaders import YoutubeLoader

import whisper
import openpyxl

# Check OpenAI key
if os.getenv("OPENAI_API_KEY"):
    print("✅ Detected OPENAI_API_KEY")
else:
    print("⚠️ Missing OPENAI_API_KEY – LLM may fail")

# Tools definitions

api_wrapper = WikipediaAPIWrapper(top_k_results=1, doc_content_chars_max=1000)
def wikipedia_search(query: str) -> str:
    return WikipediaQueryRun(api_wrapper=api_wrapper).run({"query": query})

def run_python_with_output(code: str) -> str:
    if "print(" not in code:
        code = f"print({code})"
    return PythonREPLTool().run(code)

def get_youtube_transcript(url: str) -> str:
    try:
        loader = YoutubeLoader.from_youtube_url(url, add_video_info=False)
        docs = loader.load()
        return " ".join(d.page_content for d in docs)
    except Exception as e:
        return "[YOUTUBE ERROR] " + str(e)

def transcribe_audio(file_path: str) -> str:
    try:
        model = whisper.load_model("base")
        res = model.transcribe(file_path)
        return res["text"]
    except Exception as e:
        return "[AUDIO ERROR] " + str(e)

def extract_excel_total_food_sales(file_path: str) -> str:
    try:
        wb = openpyxl.load_workbook(file_path)
        sheet = wb.active
        total = 0.0
        for _, category, amount in sheet.iter_rows(min_row=2, values_only=True):
            if isinstance(category, str) and "food" in category.lower():
                total += float(amount or 0)
        return f"${total:.2f}"
    except Exception as e:
        return "[EXCEL ERROR] " + str(e)

# Assemble tools
TOOLS = [
    FunctionTool.from_defaults(wikipedia_search),
    FunctionTool.from_defaults(run_python_with_output),
    FunctionTool.from_defaults(get_youtube_transcript),
    FunctionTool.from_defaults(transcribe_audio),
    FunctionTool.from_defaults(extract_excel_total_food_sales),
]

# LLM and Agent
llm = OpenAI(model="gpt-4")
agent = ReActAgent.from_tools(
    tools=TOOLS,
    llm=llm,
    verbose=True,
    system_prompt="""
You are an expert AI assistant on the GAIA benchmark.

Use available tools (Wikipedia, Python, YouTube transcript, audio, Excel).
Output ONLY the final answer. No reasoning or commentary.
Format exactly as requested (list, number, name, chess move, currency).
If tool fails, output "Tool not available".
""",
)

def answer_question_sync(question: str) -> str:
    try:
        resp = agent.chat(question)
        if hasattr(resp, "response") and hasattr(resp.response, "content"):
            return resp.response.content.strip()
        return str(resp).strip()
    except Exception as e:
        print("❌ Agent exception:", e)
        return "[ERROR] " + str(e)

async def answer_question(question: str) -> str:
    return answer_question_sync(question)