import gradio as gr
from api import get_all_questions, submit_answers, get_file
from agent import answer_question
from tools.file_loader import read_pdf, read_csv, read_txt
import traceback

qa_data = []
answers = []

def run_agent():
    global qa_data, answers
    try:
        print("[INFO] Fetching questions from GAIA API...")
        qa_data = get_all_questions()
        answers = []

        print(f"[INFO] Total questions retrieved: {len(qa_data)}")

        for q in qa_data:
            task_id = q.get('task_id')
            question = q.get('question')
            file_text = None

            if q.get("files"):
                for fname in q["files"]:
                    try:
                        raw = get_file(task_id, fname)
                        print(f"[INFO] Loaded file: {fname}")

                        if fname.endswith(".pdf"):
                            file_text = read_pdf(raw)
                        elif fname.endswith(".csv"):
                            file_text = read_csv(raw)
                        elif fname.endswith(".txt"):
                            file_text = read_txt(raw)
                        else:
                            print(f"[WARN] Unsupported file type: {fname}")
                    except Exception as fe:
                        print(f"[ERROR] Failed to load file {fname}: {str(fe)}")
                    break  # Load only the first file

            print(f"[INFO] Answering question: {question}")
            response = answer_question(question, file_context=file_text, do_search=True)
            print(f"[INFO] Answer: {response}")

            answers.append({
                "task_id": task_id,
                "submitted_answer": response
            })

        return f"✅ Agent answered {len(answers)} questions. Ready to submit!"

    except Exception as e:
        error_log = traceback.format_exc()
        print("[ERROR]", error_log)
        return f"❌ Error:\n{error_log}"

def handle_submit(username, code_link):
    try:
        print(f"[INFO] Submitting answers for user: {username}")
        result = submit_answers(username, code_link, answers)
        print(f"[INFO] Submission response: {result}")
        return result
    except Exception as e:
        error_log = traceback.format_exc()
        print("[ERROR]", error_log)
        return f" Submission failed:\n{error_log}"

with gr.Blocks() as demo:
    gr.Markdown("# GAIA Level 1 Agent")

    with gr.Row():
        username = gr.Textbox(label="FD900")
        code_link = gr.Textbox(label="Space Code URL (https://huggingface.co/spaces/FD900/Final_Assignment_Template/tree/main)")

    status = gr.Textbox(label="Status", lines=4)

    run_btn = gr.Button("Run Agent")
    submit_btn = gr.Button(" Submit Answers")

    run_btn.click(fn=run_agent, outputs=status)
    submit_btn.click(fn=handle_submit, inputs=[username, code_link], outputs=status)

# Make sure this launches in Spaces
demo.launch()