import gradio as gr
import os
from postpartum_agent import PostpartumResearchAgent

# Hardcoded GAIA-style test questions
gaia_tasks = [
    {"task_id": "q1", "question": "List all vegetables from: milk, eggs, flour, whole bean coffee, Oreos, sweet potatoes, fresh basil, plums, green beans, rice, corn, bell pepper, whole allspice, acorns, broccoli, celery, zucchini, lettuce, peanuts."},
    {"task_id": "q2", "question": "What is postpartum depression and how should it be addressed?"},
    {"task_id": "q3", "question": "Give a short tip on how to handle postpartum fatigue."},
    {"task_id": "q4", "question": "Name a safe exercise for postpartum recovery."},
    {"task_id": "q5", "question": "What should a mom eat to regain energy after childbirth?"}
]

# Instantiate your agent
agent = PostpartumResearchAgent()

def answer_question(question):
    return agent.run(question)

def submit_to_gaia(profile: gr.OAuthProfile):
    if profile is None:
        return "❌ Please log in to Hugging Face first!"
    username = profile.username
    # rest of your logic...
    space_id = os.getenv("SPACE_ID") or "your-username/your-space"
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"

    answers = []
    for q in gaia_tasks:
        task_id = q["task_id"]
        question_text = q["question"]
        result = agent.run(question_text)
        cleaned = (result or "").strip()
        answers.append({
            "task_id": task_id,
            "submitted_answer": cleaned
        })

    payload = {
        "username": username,
        "agent_code": agent_code,
        "answers": answers
    }

    # For testing, just show the payload instead of POSTing
    preview = "\n".join([f"{a['task_id']}: {a['submitted_answer']}" for a in answers])

    return f"""✅ Submission complete!
Username: {username}
Tasks answered: {len(answers)}
Preview of answers:
{preview}
Agent code URL: {agent_code}
"""

with gr.Blocks() as demo:
    gr.Markdown("# 🤱 Postpartum Research Agent with GAIA Submission")

    with gr.Row():
        with gr.Column():
            question = gr.Textbox(label="Ask a Postpartum Question")
            answer = gr.Textbox(label="Agent Answer", lines=6)
            ask_btn = gr.Button("Get Answer")

        with gr.Column():
            gr.Markdown("### 🌍 Submit Your Agent to GAIA Leaderboard")
            gr.LoginButton()  # ✅ Adds login button to the UI
            submit_output = gr.Textbox(label="Submission Result", lines=10)
            submit_btn = gr.Button("Submit to GAIA")

    ask_btn.click(fn=answer_question, inputs=question, outputs=answer)
    submit_btn.click(fn=submit_to_gaia, outputs=submit_output)

demo.launch()