import gradio as gr
from datasets import load_dataset, Dataset
from datetime import datetime, date 
import io
import os
from PIL import Image, ImageDraw, ImageFont
from huggingface_hub import login
import requests 
import json 
import base64 # <-- ADDED IMPORT for image handling

# Attempt to login using environment token
try:
    HF_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")
    if HF_TOKEN:
        login(token=HF_TOKEN)
        print("Logged in to Hugging Face Hub successfully.")
    else:
        print("HUGGINGFACE_TOKEN environment variable not set.")
except Exception as e:
    print(f"Error logging in to Hugging Face Hub: {e}")

# Constants for Certificate Generation
SCORES_DATASET = "agents-course/unit4-students-scores"
CERTIFICATES_DATASET = "agents-course/course-certificates-of-excellence"
THRESHOLD_SCORE = 30

# --- Constants for GAIA Benchmark API ---
GAIA_API_BASE_URL = "https://agents-course-unit4-scoring.hf.space"

# --- Constants for Gemini API ---
GEMINI_API_URL_BASE = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent"

# --- Functions for Certificate Generation (existing code) ---
def check_user_score(username):
    try:
        score_data = load_dataset(SCORES_DATASET, split="train", download_mode="force_redownload", token=HF_TOKEN if HF_TOKEN else True)
        matches = [row for row in score_data if row["username"] == username]
        return matches[0] if matches else None
    except Exception as e:
        print(f"Error checking user score: {e}")
        return None

def has_certificate_entry(username):
    try:
        cert_data = load_dataset(CERTIFICATES_DATASET, split="train", download_mode="force_redownload", token=HF_TOKEN if HF_TOKEN else True)
        return any(row["username"] == username for row in cert_data)
    except Exception as e:
        print(f"Error checking certificate entry: {e}")
        return False

def add_certificate_entry(username, name, score):
    try:
        ds = load_dataset(CERTIFICATES_DATASET, split="train", download_mode="force_redownload", token=HF_TOKEN if HF_TOKEN else True)
        filtered_rows = [row for row in ds if row["username"] != username]
        new_entry = {
            "username": username,
            "score": score,
            "name_on_certificate": name,
            "timestamp": datetime.now().isoformat()
        }
        filtered_rows.append(new_entry)
        updated_ds = Dataset.from_list(filtered_rows)
        updated_ds.push_to_hub(CERTIFICATES_DATASET, token=HF_TOKEN if HF_TOKEN else None)
        print(f"Certificate entry added/updated for {username}.")
    except Exception as e:
        print(f"Error adding certificate entry: {e}")

def generate_certificate_image(name_on_cert):
    try:
        current_dir = os.path.dirname(__file__)
        certificate_template_path = os.path.join(current_dir, "certificate.png")
        font_path = os.path.join(current_dir, "Quattrocento-Regular.ttf")

        if not os.path.exists(certificate_template_path):
            alt_cert_path_templates_parent = os.path.join(current_dir,"..", "templates", "certificate.png")
            alt_cert_path_root = os.path.join(current_dir, "certificate.png") 
            
            if os.path.exists(alt_cert_path_templates_parent):
                 certificate_template_path = alt_cert_path_templates_parent
            elif os.path.exists(alt_cert_path_root):
                 certificate_template_path = alt_cert_path_root
            else:
                raise FileNotFoundError(f"Certificate template not found. Checked default, ../templates/, and root relative to app.py.")

        if not os.path.exists(font_path):
            alt_font_path_parent = os.path.join(current_dir, "..","Quattrocento-Regular.ttf")
            alt_font_path_root = os.path.join(current_dir, "Quattrocento-Regular.ttf") 

            if os.path.exists(alt_font_path_parent):
                 font_path = alt_font_path_parent
            elif os.path.exists(alt_font_path_root):
                 font_path = alt_font_path_root
            else:
                raise FileNotFoundError(f"Font file not found. Checked default and parent directory relative to app.py.")

        im = Image.open(certificate_template_path)
        d = ImageDraw.Draw(im)
        name_font = ImageFont.truetype(font_path, 100)
        date_font = ImageFont.truetype(font_path, 48)
        name_on_cert = name_on_cert.title()
        d.text((1000, 740), name_on_cert, fill="black", anchor="mm", font=name_font)
        d.text((1480, 1170), str(date.today()), fill="black", anchor="mm", font=date_font)
        pdf_buffer = io.BytesIO()
        im.convert("RGB").save(pdf_buffer, format="PDF")
        pdf_buffer.seek(0)
        return im, pdf_buffer
    except FileNotFoundError as fnf_error:
        print(fnf_error)
        raise 
    except Exception as e:
        print(f"Error generating certificate image: {e}")
        raise

def handle_certificate(name_on_certificate_input, profile: gr.OAuthProfile):
    if not profile: 
        return "You must be logged in with your Hugging Face account.", None, None
    username = profile.username
    if not name_on_certificate_input.strip():
        return "Please enter the name you want on the certificate.", None, None
    user_score_info = check_user_score(username)
    if not user_score_info:
        return f"No score found for {username}. Please complete Unit 4 first by submitting your agent's answers.", None, None
    score = user_score_info.get("score", 0) 
    if score < THRESHOLD_SCORE:
        return f"Your score is {score}. You need at least {THRESHOLD_SCORE} to pass and receive a certificate.", None, None
    try:
        certificate_image, pdf_bytesio_object = generate_certificate_image(name_on_certificate_input)
        add_certificate_entry(username, name_on_certificate_input, score)
        temp_pdf_path = f"certificate_{username}.pdf"
        with open(temp_pdf_path, "wb") as f:
            f.write(pdf_bytesio_object.getvalue())
        return f"Congratulations, {name_on_certificate_input}! You scored {score}. Here's your certificate:", certificate_image, temp_pdf_path
    except FileNotFoundError as e:
        return f"Critical error: A required file for certificate generation was not found: {e}. Please check Space file structure.", None, None
    except Exception as e:
        print(f"An unexpected error occurred in handle_certificate: {e}")
        return "An unexpected error occurred while generating your certificate. Please try again later.", None, None

# --- Functions for GAIA Benchmark Interaction ---
def get_gaia_api_questions():
    try:
        questions_url = f"{GAIA_API_BASE_URL}/questions"
        print(f"Attempting to fetch questions from: {questions_url}")
        response = requests.get(questions_url, timeout=30) 
        response.raise_for_status() 
        return response.json(), None 
    except requests.exceptions.RequestException as e:
        print(f"Error fetching GAIA questions: {e}")
        return None, f"Error fetching GAIA questions: {e}"
    except Exception as e:
        print(f"An unexpected error occurred while fetching questions: {e}")
        return None, f"An unexpected error occurred: {e}"

def get_gaia_file_data_for_task(task_id_for_file_fetch, associated_file_metadata_list):
    """
    Fetches the content of the primary file associated with a task_id from the GAIA API.
    Returns raw_bytes, detected_mime_type, and file_name.
    associated_file_metadata_list is the 'files' list from the question data.
    """
    # If no metadata, assume no file to fetch for this specialized getter.
    # Or, if the API always serves THE file for task_id, then metadata is just for info.
    # Let's assume the API /files/{task_id} always gives the relevant file if one exists for the task.
    
    file_url = f"{GAIA_API_BASE_URL}/files/{task_id_for_file_fetch}"
    print(f"Attempting to fetch file for task {task_id_for_file_fetch} from {file_url}")
    
    try:
        response = requests.get(file_url, timeout=30)
        response.raise_for_status() # This will error if file not found (404) or other issues
        
        raw_bytes = response.content
        detected_mime_type = response.headers.get('Content-Type', '').split(';')[0].strip()
        
        # Try to get a filename from metadata if available, otherwise default
        file_name = "attached_file"
        if associated_file_metadata_list and isinstance(associated_file_metadata_list, list) and len(associated_file_metadata_list) > 0:
            # Assuming the first file in metadata is the one fetched, or provides its name
            first_file_meta = associated_file_metadata_list[0]
            if isinstance(first_file_meta, dict) and 'file_name' in first_file_meta:
                file_name = first_file_meta['file_name']
        
        print(f"File fetched for task {task_id_for_file_fetch}. Mime-type: {detected_mime_type}, Name: {file_name}, Size: {len(raw_bytes)} bytes")
        return raw_bytes, detected_mime_type, file_name
        
    except requests.exceptions.HTTPError as http_err:
        # Specifically handle 404 for "no file" vs other errors
        if http_err.response.status_code == 404:
            print(f"No file found (404) for task {task_id_for_file_fetch} at {file_url}.")
        else:
            print(f"HTTP error fetching file for task {task_id_for_file_fetch}: {http_err}")
        return None, None, None
    except requests.exceptions.RequestException as e:
        print(f"Could not fetch file for task {task_id_for_file_fetch}: {e}. Proceeding without file content.")
        return None, None, None
    except Exception as e_gen:
        print(f"Unexpected error fetching file for task {task_id_for_file_fetch}: {e_gen}")
        return None, None, None

def my_agent_logic(task_id: str, question: str, files_metadata: list = None): # files_metadata is the list from task.get("files")
    """
    Uses the Gemini API, with GAIA-specific prompting and basic file handling, 
    to generate an answer for the given question.
    """
    print(f"Agent (GAIA-Grounded Gemini) processing Task ID: {task_id}, Question: {question}")
    if files_metadata: # This is the list of file metadata dicts
        print(f"File metadata associated with this task: {files_metadata}")

    gemini_api_key = os.environ.get("GEMINI_API_KEY")
    if not gemini_api_key:
        print("Error: GEMINI_API_KEY not found in environment variables. Please set it in Space Secrets.")
        return f"ERROR_GEMINI_KEY_MISSING_FOR_TASK_{task_id}"

    # --- GAIA-specific System Prompt ---
    # Adapted from Figure 2 of GAIA Paper [cite: 103, 104, 105, 106, 107, 108]
    system_prompt_lines = [
        "You are a general AI assistant. I will ask you a question.",
        "Report your thoughts (for your own processing, not for the final answer), and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER].", # Instructing the LLM about the template it should "think" in
        "However, your actual returned response to me (the user) should ONLY be [YOUR FINAL ANSWER] part, without the 'FINAL ANSWER:' prefix.", # Clarification for our use case
        "YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.",
        "If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.",
        "If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.",
        "If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.",
        "Be precise and ensure the answer strictly adheres to any format requested in the question.",
        "If external files are mentioned or provided, use their content if relevant and accessible to answer the question.",
    ]
    # We won't send this as a separate "system" message in Gemini's typical API structure,
    # but rather prepend it to the user question for a single turn.
    
    # --- Prepare parts for Gemini API payload ---
    gemini_parts = []

    # Prepend system prompt guidelines to the main question text part
    user_question_text = "\n".join(system_prompt_lines) + f"\n\nGAIA Question: {question}"
    
    # --- File Handling ---
    file_content_bytes, detected_mime_type, file_name = None, None, None
    if files_metadata: # If the question has associated file(s) metadata
        file_content_bytes, detected_mime_type, file_name = get_gaia_file_data_for_task(task_id, files_metadata)

    if file_content_bytes:
        if detected_mime_type and detected_mime_type.startswith("image/"): # Handle images
            try:
                base64_image = base64.b64encode(file_content_bytes).decode('utf-8')
                gemini_parts.append({"text": user_question_text}) # Question text first
                gemini_parts.append({
                    "inline_data": {
                        "mime_type": detected_mime_type,
                        "data": base64_image
                    }
                })
                print(f"Added image {file_name} ({detected_mime_type}) to Gemini prompt for task {task_id}.")
            except Exception as e_img:
                print(f"Error processing image file {file_name} for task {task_id}: {e_img}")
                gemini_parts.append({"text": user_question_text + f"\n[Agent note: An image file '{file_name}' was associated but could not be processed: {e_img}]"})
        elif detected_mime_type and detected_mime_type == "text/plain": # Handle plain text files
            try:
                text_content = file_content_bytes.decode('utf-8')
                user_question_text += f"\n\nContent of attached text file '{file_name}':\n{text_content}"
                gemini_parts.append({"text": user_question_text})
                print(f"Added text file content '{file_name}' to Gemini prompt for task {task_id}.")
            except Exception as e_txt:
                print(f"Error decoding text file {file_name} for task {task_id}: {e_txt}")
                gemini_parts.append({"text": user_question_text + f"\n[Agent note: A text file '{file_name}' was associated but could not be decoded: {e_txt}]"})
        else: # Other file types, just mention them
            user_question_text += f"\n\nNote: A file named '{file_name}' (type: {detected_mime_type or 'unknown'}) is associated with this question. Its content is not directly viewable in this text prompt."
            gemini_parts.append({"text": user_question_text})
            print(f"Noted non-image/text file {file_name} ({detected_mime_type}) in Gemini prompt for task {task_id}.")
    else: # No file content fetched or no files associated
        gemini_parts.append({"text": user_question_text})

    payload = {
        "contents": [{"role": "user", "parts": gemini_parts}],
        "generationConfig": { 
            "temperature": 0.2, # Lower temperature for more factual/deterministic GAIA answers
            "maxOutputTokens": 300, # Increased slightly for potentially more complex answers
        }
    }

    api_url_with_key = f"{GEMINI_API_URL_BASE}?key={gemini_api_key}"
    agent_computed_answer = f"ERROR_CALLING_GEMINI_FOR_TASK_{task_id}" 

    try:
        headers = {"Content-Type": "application/json"}
        print(f"Calling Gemini API for task {task_id}...")
        response = requests.post(api_url_with_key, headers=headers, json=payload, timeout=60) 
        response.raise_for_status()  
        result = response.json()
        
        if (result.get("candidates") and 
            result["candidates"][0].get("content") and
            result["candidates"][0]["content"].get("parts") and
            result["candidates"][0]["content"]["parts"][0].get("text")):
            raw_answer = result["candidates"][0]["content"]["parts"][0]["text"].strip()
            
            # Remove the "FINAL ANSWER:" prefix if the LLM included it, despite instructions
            if raw_answer.upper().startswith("FINAL ANSWER:"):
                agent_computed_answer = raw_answer[len("FINAL ANSWER:"):].strip()
            else:
                agent_computed_answer = raw_answer
            # Further cleaning: sometimes LLMs might still add subtle quotes if the answer is a simple string
            if len(agent_computed_answer) > 1 and ((agent_computed_answer.startswith('"') and agent_computed_answer.endswith('"')) or \
               (agent_computed_answer.startswith("'") and agent_computed_answer.endswith("'"))):
                agent_computed_answer = agent_computed_answer[1:-1]

        else:
            print(f"Warning: Unexpected response structure from Gemini API for task {task_id}: {result}")
            if result.get("promptFeedback") and result["promptFeedback"].get("blockReason"):
                block_reason = result["promptFeedback"]["blockReason"]
                print(f"Gemini API blocked the prompt for task {task_id}. Reason: {block_reason}")
                agent_computed_answer = f"ERROR_GEMINI_PROMPT_BLOCKED_{block_reason}_FOR_TASK_{task_id}"
            else:
                agent_computed_answer = f"ERROR_PARSING_GEMINI_RESPONSE_FOR_TASK_{task_id}"
    except requests.exceptions.Timeout:
        print(f"Timeout error calling Gemini API for task {task_id}.")
        agent_computed_answer = f"ERROR_GEMINI_TIMEOUT_FOR_TASK_{task_id}"
    except requests.exceptions.RequestException as e:
        print(f"Error calling Gemini API for task {task_id}: {e}")
        if e.response is not None:
            print(f"Gemini API Error Response Status: {e.response.status_code}")
            try: print(f"Gemini API Error Response Body: {e.response.json()}")
            except json.JSONDecodeError: print(f"Gemini API Error Response Body (text): {e.response.text}")
        agent_computed_answer = f"ERROR_GEMINI_REQUEST_FAILED_FOR_TASK_{task_id}"
    except Exception as e:
        print(f"An unexpected error occurred in my_agent_logic for task {task_id}: {e}")
        agent_computed_answer = f"ERROR_UNEXPECTED_IN_AGENT_LOGIC_FOR_TASK_{task_id}"

    print(f"Agent (GAIA-Grounded Gemini) computed answer for Task ID {task_id}: {agent_computed_answer}")
    return agent_computed_answer

def run_agent_on_gaia(profile: gr.OAuthProfile, run_all_questions: bool = True):
    if not profile:
        return "You must be logged in to run the agent.", None, None
    log_messages = ["Starting agent run..."]
    questions_data, error_msg = get_gaia_api_questions()
    if error_msg:
        log_messages.append(error_msg)
        return "\n".join(log_messages), None, None
    if not questions_data:
        log_messages.append("No questions retrieved from GAIA API.")
        return "\n".join(log_messages), None, None
    log_messages.append(f"Retrieved {len(questions_data)} questions from GAIA.")
    answers_to_submit = []
    tasks_to_process = questions_data
    if not run_all_questions and questions_data:
        import random
        if not isinstance(questions_data, list) or not questions_data:
             log_messages.append("Question data is not a list or is empty, cannot pick random.")
             return "\n".join(log_messages), None, None
        tasks_to_process = [random.choice(questions_data)]
        log_messages.append(f"Processing 1 random question based on user choice.")
    elif run_all_questions:
        log_messages.append(f"Processing all {len(tasks_to_process)} questions.")
    for task in tasks_to_process:
        task_id = task.get("task_id")
        question = task.get("question")
        associated_files_metadata = task.get("files", []) # This is the list of file metadata dicts
        if task_id and question:
            log_messages.append(f"\nProcessing Task ID: {task_id}")
            log_messages.append(f"Question: {question}")
            if associated_files_metadata:
                 log_messages.append(f"Associated files metadata: {associated_files_metadata}")
            # Pass the files_metadata to the agent logic
            submitted_answer = my_agent_logic(task_id, question, associated_files_metadata)
            log_messages.append(f"Agent's Answer: {submitted_answer}")
            answers_to_submit.append({"task_id": task_id, "submitted_answer": submitted_answer})
        else:
            log_messages.append(f"Skipping malformed task: {task}")
    if not answers_to_submit:
        log_messages.append("No answers were generated by the agent.")
    return "\n".join(log_messages), answers_to_submit, answers_to_submit

def submit_agent_answers(profile: gr.OAuthProfile, answers_for_submission_state):
    if not profile:
        return "You must be logged in to submit answers."
    if not answers_for_submission_state:
        return "No answers available to submit. Please run the agent first."
    username = profile.username
    space_id = os.getenv('SPACE_ID', '') 
    agent_code_link = f"https://huggingface.co/spaces/{space_id}/tree/main"
    submission_log_messages = [f"Preparing to submit answers for user: {username}"]
    if not space_id:
        your_space_name_guess = os.path.basename(os.path.dirname(os.path.abspath(__file__)))
        if not your_space_name_guess or your_space_name_guess == 'app':
             your_space_name_guess = "YOUR_SPACE_NAME_HERE" 
        agent_code_link = f"https://huggingface.co/spaces/{username}/{your_space_name_guess}/tree/main"
        submission_log_messages.append(f"Warning: SPACE_ID not found. Constructed agent_code_link as: {agent_code_link}. Please verify this link is correct.")
    submission_log_messages.append(f"Agent Code Link: {agent_code_link}")
    payload = {
        "username": username,
        "agent_code": agent_code_link,
        "answers": answers_for_submission_state 
    }
    try:
        submit_url = f"{GAIA_API_BASE_URL}/submit"
        print(f"Attempting to submit answers to: {submit_url} with payload: {payload}") 
        response = requests.post(submit_url, json=payload, timeout=60) 
        response.raise_for_status()
        submission_response = response.json()
        submission_log_messages.append(f"Submission successful! Response: {submission_response}")
        message = submission_response.get("message")
        score = submission_response.get("score")
        score_string = submission_response.get("score_string")
        final_message = "Submission processed."
        if message: final_message = message
        elif score_string: final_message = score_string
        elif score is not None: final_message = f"Score: {score}"
        return "\n".join(submission_log_messages) + f"\n\n➡️ Result: {final_message}"
    except requests.exceptions.Timeout:
        error_detail = f"Timeout error submitting answers to GAIA scoring API."
        submission_log_messages.append(error_detail)
        return "\n".join(submission_log_messages)
    except requests.exceptions.RequestException as e:
        error_detail = f"Error submitting answers: {e}"
        if e.response is not None:
            error_detail += f"\nResponse status: {e.response.status_code}"
            try: error_detail += f"\nResponse body: {e.response.json()}"
            except ValueError: error_detail += f"\nResponse body (text): {e.response.text}"
        submission_log_messages.append(error_detail)
        return "\n".join(submission_log_messages)
    except Exception as e:
        submission_log_messages.append(f"An unexpected error occurred during submission: {e}")
        return "\n".join(submission_log_messages)

# --- Gradio Interface (largely unchanged from your latest version) ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎓 Agents Course - Unit 4 Final Project")
    gr.Markdown("⚠️ **Note**: Due to high demand, you might experience occasional bugs. If something doesn't work, please try again after a moment!")
    gr.Markdown("---")
    gr.Markdown("Your Hugging Face login token (`HUGGINGFACE_TOKEN`) should be set as a Space Secret for dataset pushes.")
    gr.Markdown("Your Gemini API Key (`GEMINI_API_KEY`) **MUST** be set as a Space Secret for the agent to function.")
    gr.Markdown(f"**GAIA API Base URL Used:** `{GAIA_API_BASE_URL}`")
    gr.LoginButton() 
    answers_to_submit_state = gr.State([])
    with gr.Tabs():
        with gr.TabItem("🤖 Run Agent on GAIA Benchmark"):
            gr.Markdown("## Step 1: Run Your Agent & Generate Answers")
            gr.Markdown("This agent uses the Gemini API (with GAIA-specific prompting and basic file handling) to generate answers.")
            run_all_questions_checkbox = gr.Checkbox(label="Process all questions (unchecked processes 1 random question for testing)", value=True)
            run_agent_button = gr.Button("🔎 Fetch Questions & Run My Agent")
            gr.Markdown("### Agent Run Log & Generated Answers:")
            agent_run_log_display = gr.Textbox(label="Agent Run Log", lines=10, interactive=False)
            generated_answers_display = gr.JSON(label="Generated Answers (for review before submission)") 
            run_agent_button.click(
                fn=run_agent_on_gaia,
                inputs=[run_all_questions_checkbox], 
                outputs=[agent_run_log_display, generated_answers_display, answers_to_submit_state] 
            )
            gr.Markdown("## Step 2: Submit Agent's Answers")
            gr.Markdown("Once you have reviewed the generated answers, click below to submit them for scoring.")
            submit_button = gr.Button("🚀 Submit Answers to GAIA Benchmark")
            submission_status_display = gr.Textbox(label="Submission Status", lines=5, interactive=False)
            submit_button.click(
                fn=submit_agent_answers,
                inputs=[answers_to_submit_state], 
                outputs=[submission_status_display]
            )
        with gr.TabItem("🏅 Get Certificate"):
            gr.Markdown("# ✅ How to Get Your Certificate (After Scoring >= 30%)")
            gr.Markdown("""
            1. Ensure you are logged in.
            2. If your agent scored 30% or higher (after submitting on the 'Run Agent' tab), you can get your certificate.
            3. Enter your full name as you want it to appear on the certificate.
            4. Click 'Get My Certificate'.
            """)
            gr.Markdown("---")
            gr.Markdown("📝 **Note**: You must have successfully submitted your agent's answers and achieved a score of **30% or higher**.")
            with gr.Row():
                name_input = gr.Text(label="Enter your full name (this will appear on the certificate)")
            generate_cert_btn = gr.Button("📜 Get My Certificate")
            output_text_cert = gr.Textbox(label="Certificate Result")
            cert_image_display = gr.Image(label="Your Certificate Image", type="pil") 
            cert_file_download = gr.File(label="Download Certificate (PDF)") 
            generate_cert_btn.click(
                fn=handle_certificate,
                inputs=[name_input], 
                outputs=[output_text_cert, cert_image_display, cert_file_download]
            )

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