import aiohttp
import asyncio
from graph import graph
from state import JARVISState
from pydantic import BaseModel
from typing import List
import json
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
# Debug: Verify environment variables
print(f"OPENAI_API_KEY loaded: {'set' if os.getenv('OPENAI_API_KEY') else 'not set'}")
print(f"LANGFUSE_PUBLIC_KEY loaded: {'set' if os.getenv('LANGFUSE_PUBLIC_KEY') else 'not set'}")

# Verify critical environment variables
required_env_vars = ["OPENAI_API_KEY", "LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY"]
for var in required_env_vars:
    if not os.getenv(var):
        raise ValueError(f"Environment variable {var} is not set")

# Pydantic Models for Submission
class Answer(BaseModel):
    task_id: str
    submitted_answer: str

class Submission(BaseModel):
    username: str
    agent_code: str
    answers: List[Answer]

async def fetch_questions() -> List[dict]:
    async with aiohttp.ClientSession() as session:
        async with session.get("https://api.gaia-benchmark.com/questions") as resp:
            return await resp.json()

async def download_file(task_id: str, file_path: str) -> bool:
    async with aiohttp.ClientSession() as session:
        async with session.get(f"https://api.gaia-benchmark.com/files/{task_id}") as resp:
            if resp.status == 200:
                with open(file_path, "wb") as f:
                    f.write(await resp.read())
                return True
            return False

async def process_question(question: dict) -> Answer:
    # Determine file type based on question context
    file_type = "jpg" if "image" in question["question"].lower() else "txt"
    if "menu" in question["question"].lower() or "report" in question["question"].lower() or "document" in question["question"].lower():
        file_type = "pdf"  # Prioritize PDF for reports/documents
    elif "data" in question["question"].lower():
        file_type = "csv"
    
    file_path = f"temp_{question['task_id']}.{file_type}"
    await download_file(question["task_id"], file_path)
    
    state = JARVISState(
        task_id=question["task_id"],
        question=question["question"],
        tools_needed=[],
        web_results=[],
        file_results="",
        image_results="",
        calculation_results="",
        document_results="",
        messages=[],
        answer=""
    )
    # Use unique thread_id for memory
    result = await graph.ainvoke(state, config={"thread_id": question["task_id"]})
    return Answer(task_id=question["task_id"], submitted_answer=result["answer"])

async def submit_answers(answers: List[Answer], username: str, agent_code: str):
    submission = Submission(
        username=username,
        agent_code=agent_code,
        answers=answers
    )
    async with aiohttp.ClientSession() as session:
        async with session.post("https://api.gaia-benchmark.com/submit", json=submission.dict()) as resp:
            return await resp.json()

async def main():
    username = "onisj"  # Your Hugging Face username
    agent_code = "https://huggingface.co/spaces/onisj/jarvis_gaia_agent/tree/main"
    questions = await fetch_questions()
    answers = []
    for question in questions[:20]:  # Process 20 questions
        answer = await process_question(question)
        answers.append(answer)
    result = await submit_answers(answers, username, agent_code)
    print("Submission result:", json.dumps(result, indent=2))

if __name__ == "__main__":
    asyncio.run(main())