import os
from openai import OpenAI

class GaiaAgent:
    def __init__(self):
        self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
        self.instructions = (
            "You are a research assistant solving GAIA benchmark questions using 2022 English Wikipedia knowledge.\n"
            "For each question, reason step-by-step and only return the final answer in exact format."
        )
        self.task_templates = {
            "8e867cd7-cff9-4e6c-867a-ff5ddc2550be": self.q_mercedes_sosa_albums
        }

    def __call__(self, question: str, task_id: str = None) -> str:
        if task_id in self.task_templates:
            return self.task_templates[task_id](question)
        else:
            return "[SKIPPED: Task not yet implemented in Agent V9]"

    def q_mercedes_sosa_albums(self, question: str) -> str:
        prompt = (
            "You are a research assistant using 2022 English Wikipedia knowledge.\n"
            "\nQUESTION:\n"
            "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)?\n"
            "\nFollow this step-by-step reasoning chain:\n"
            "Step 1: Identify all studio albums by Mercedes Sosa.\n"
            "Step 2: Filter the albums published between 2000 and 2009 inclusive.\n"
            "Step 3: Count them.\n"
            "Step 4: Return only the count as a number (no text, no explanation).\n"
            "\nANSWER:"
        )

        try:
            response = self.client.chat.completions.create(
                model="gpt-4-turbo",
                messages=[
                    {"role": "system", "content": self.instructions},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.0
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            return f"[AGENT ERROR: {e}]"