import os
import requests
import pandas as pd
import gradio as gr
import spaces

from io import BytesIO
from smolagents import ToolCallingAgent, DuckDuckGoSearchTool, OpenAIServerModel

API_URL = "https://agents-course-unit4-scoring.hf.space"
AGENT_CODE_URL = "https://huggingface.co/spaces/bacchi87/Final_Assignment_Template/tree/main"

# Gemini model setup
model = OpenAIServerModel(
    model_id="gemini-3.5-flash",
    api_base="https://generativelanguage.googleapis.com/v1beta/openai/",
    api_key=os.environ["GEMINI_API_KEY"],
    max_tokens=1500
)

agent = ToolCallingAgent(
    tools=[DuckDuckGoSearchTool()],
    model=model,
    additional_authorized_imports=[
        "math", "re", "json", "datetime", "pandas"
    ],
    max_steps=6
)


def test_gemini():
    try:
        result = agent.run(
            "Reply with exactly GEMINI_OK. Do not use tools."
        )
        return f"Gemini connection successful!\nResponse: {result}"
    except Exception as e:
        return f"Gemini connection failed:\n{e}"


def get_attachment(task_id):
    try:
        response = requests.get(
            f"{API_URL}/files/{task_id}", timeout=30
        )
        response.raise_for_status()
        data = BytesIO(response.content)

        try:
            df = pd.read_excel(data)
            return df.to_string(index=False)
        except Exception:
            data.seek(0)
            try:
                df = pd.read_csv(data)
                return df.to_string(index=False)
            except Exception:
                return "Attachment format is not readable."

    except Exception as e:
        return f"Attachment download failed: {e}"


def answer_question(question, attachment=""):
    prompt = f"""
Answer this question accurately.

Question:
{question}

Attached file data:
{attachment}

Use web search or calculations when required.
Return ONLY the answer.
Do not write FINAL ANSWER.
"""
    try:
        result = agent.run(prompt)
        return str(result).strip()
    except Exception as e:
        return f"Agent error: {e}"


@spaces.GPU
def run_evaluation(profile: gr.OAuthProfile | None):

    if profile is None:
        return "Please login with Hugging Face.", []

    username = profile.username
    log = f"Logged in as: {username}\n"

    try:
        response = requests.get(
            f"{API_URL}/questions", timeout=30
        )
        response.raise_for_status()
        questions = response.json()

        if not questions:
            return "No questions returned.", []

    except Exception as e:
        return f"Could not fetch questions: {e}", []

    answers = []

    for i, item in enumerate(questions, start=1):

        task_id = str(item["task_id"])
        question = item["question"]

        log += f"\nQuestion {i}/{len(questions)}: {task_id}\n"

        attachment = ""

        if item.get("file_name"):
            attachment = get_attachment(task_id)

        answer = answer_question(question, attachment)

        # Stop if Gemini fails
        if answer.startswith("Agent error:"):
            log += f"Model inference failed: {answer}\n"
            return log + "\nNo answers submitted.", answers

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

        log += f"Answer: {answer}\n"

    # Official scoring API
    payload = {
        "username": username,
        "agent_code": AGENT_CODE_URL,
        "answers": answers
    }

    try:
        result = requests.post(
            f"{API_URL}/submit",
            json=payload,
            timeout=60
        )
        result.raise_for_status()

        score_data = result.json()

        log += "\n--- EVALUATION RESULT ---\n"
        log += f"Score: {score_data.get('score')}%\n"
        log += f"Correct: {score_data.get('correct_count')}\n"
        log += f"Attempted: {score_data.get('total_attempted')}\n"
        log += f"Message: {score_data.get('message')}\n"

    except Exception as e:
        log += f"\nSubmission failed: {e}\n"

    return log, answers


with gr.Blocks() as demo:

    gr.Markdown("# Hugging Face Unit 4 - GAIA Evaluation")

    gr.LoginButton()

    test_btn = gr.Button("Test Gemini Connection")
    test_output = gr.Textbox(label="Gemini Test Result")

    test_btn.click(
        fn=test_gemini,
        inputs=[],
        outputs=[test_output]
    )

    run_btn = gr.Button(
        "Run Evaluation & Submit Answers",
        variant="primary"
    )

    status_box = gr.Textbox(
        label="Evaluation Status",
        lines=18
    )

    answers_box = gr.JSON(label="Generated Answers")

    run_btn.click(
        fn=run_evaluation,
        inputs=[],
        outputs=[status_box, answers_box]
    )


if __name__ == "__main__":
    demo.launch()