import os
import tempfile
from pathlib import Path

import gradio as gr
import requests
import pandas as pd


# ============================================================
# CONFIGURATION
# ============================================================

DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"

HF_TOKEN = os.getenv("HF_TOKEN")

# Model used by the agent.
# Hugging Face's current Agents Course examples use
# InferenceClientModel with CodeAgent.
MODEL_ID = os.getenv(
    "MODEL_ID",
    "Qwen/Qwen2.5-Coder-32B-Instruct"
)


# ============================================================
# SMOLAGENTS IMPORTS
# ============================================================

from smolagents import (
    CodeAgent,
    InferenceClientModel,
    DuckDuckGoSearchTool,
    PythonInterpreterTool,
    VisitWebpageTool,
    SpeechToTextTool,
    Tool,
)


# ============================================================
# CUSTOM TOOL: READ LOCAL FILE
# ============================================================

class ReadLocalFileTool(Tool):

    name = "read_local_file"

    description = """
    Read the contents of a local text, Python, CSV, JSON,
    Markdown, or other text-based file.

    Use this tool when a GAIA question provides a local file.
    """

    inputs = {
        "file_path": {
            "type": "string",
            "description": "Path of the file to read."
        }
    }

    output_type = "string"

    def forward(self, file_path):

        try:

            path = Path(file_path)

            if not path.exists():
                return f"File does not exist: {file_path}"

            # Limit extremely large files
            text = path.read_text(
                encoding="utf-8",
                errors="ignore"
            )

            if len(text) > 50000:
                text = text[:50000]

            return text

        except Exception as e:

            return f"Could not read file: {e}"


# ============================================================
# CUSTOM TOOL: READ EXCEL
# ============================================================

class ReadExcelTool(Tool):

    name = "read_excel_file"

    description = """
    Read an Excel XLSX file using pandas.

    Returns sheet names and the contents of each sheet.
    Use this for GAIA questions involving spreadsheets,
    sales tables, numbers, or Excel data.
    """

    inputs = {
        "file_path": {
            "type": "string",
            "description": "Path to an XLSX file."
        }
    }

    output_type = "string"

    def forward(self, file_path):

        try:

            path = Path(file_path)

            if not path.exists():
                return f"Excel file does not exist: {file_path}"

            excel = pd.ExcelFile(path)

            output = []

            for sheet in excel.sheet_names:

                df = pd.read_excel(
                    path,
                    sheet_name=sheet
                )

                output.append(
                    f"\n--- SHEET: {sheet} ---\n"
                )

                output.append(
                    df.to_string(index=False)
                )

            result = "\n".join(output)

            if len(result) > 50000:
                result = result[:50000]

            return result

        except Exception as e:

            return f"Could not read Excel file: {e}"


# ============================================================
# CUSTOM TOOL: FILE INFORMATION
# ============================================================

class FileInfoTool(Tool):

    name = "file_information"

    description = """
    Inspect a local file and return its name, extension,
    size, and basic information.
    """

    inputs = {
        "file_path": {
            "type": "string",
            "description": "Path to the local file."
        }
    }

    output_type = "string"

    def forward(self, file_path):

        try:

            path = Path(file_path)

            if not path.exists():
                return "File does not exist."

            size = path.stat().st_size

            return (
                f"File name: {path.name}\n"
                f"Extension: {path.suffix}\n"
                f"Size: {size} bytes\n"
                f"Path: {path}"
            )

        except Exception as e:

            return f"Error inspecting file: {e}"


# ============================================================
# CUSTOM TOOL: YOUTUBE TRANSCRIPT
# ============================================================

class YouTubeTranscriptTool(Tool):

    name = "youtube_transcript"

    description = """
    Retrieve a transcript from a YouTube video when captions
    are available.

    Use this for questions asking what someone said in a
    YouTube video or asking about spoken dialogue.
    """

    inputs = {
        "video_url": {
            "type": "string",
            "description": "Full YouTube video URL."
        }
    }

    output_type = "string"

    def forward(self, video_url):

        try:

            from youtube_transcript_api import (
                YouTubeTranscriptApi
            )

            # Extract video ID
            video_id = None

            if "v=" in video_url:
                video_id = video_url.split("v=")[1].split("&")[0]

            elif "youtu.be/" in video_url:
                video_id = video_url.split("youtu.be/")[1].split("?")[0]

            elif "youtube.com/shorts/" in video_url:
                video_id = video_url.split("youtube.com/shorts/")[1].split("?")[0]

            if not video_id:
                return "Could not extract YouTube video ID."

            api = YouTubeTranscriptApi()

            transcript = api.fetch(video_id)

            text_parts = []

            for item in transcript:

                try:
                    text_parts.append(item.text)
                except Exception:
                    text_parts.append(str(item))

            result = " ".join(text_parts)

            if len(result) > 60000:
                result = result[:60000]

            return result

        except Exception as e:

            return (
                "Could not retrieve YouTube transcript. "
                f"Reason: {e}"
            )


# ============================================================
# GAIA AGENT
# ============================================================

class BasicAgent:

    def __init__(self):

        print("=" * 60)
        print("Initializing GAIA Agent")
        print("=" * 60)

        # ----------------------------------------------------
        # MODEL
        # ----------------------------------------------------

        model_kwargs = {
            "model_id": MODEL_ID,
            "temperature": 0.1,
            "max_tokens": 4096,
        }

        if HF_TOKEN:
            model_kwargs["token"] = HF_TOKEN

        self.model = InferenceClientModel(
            **model_kwargs
        )

        # ----------------------------------------------------
        # TOOLS
        # ----------------------------------------------------

        tools = []

        # Web search
        try:

            tools.append(
                DuckDuckGoSearchTool(
                    max_results=8,
                    rate_limit=1.0
                )
            )

            print("Web search tool enabled.")

        except Exception as e:

            print(
                f"Could not initialize web search: {e}"
            )

        # Webpage reader
        try:

            tools.append(
                VisitWebpageTool(
                    max_output_length=30000
                )
            )

            print("Webpage tool enabled.")

        except Exception as e:

            print(
                f"Could not initialize webpage tool: {e}"
            )

        # Python
        try:

            tools.append(
                PythonInterpreterTool(
                    timeout_seconds=30
                )
            )

            print("Python tool enabled.")

        except Exception as e:

            print(
                f"Could not initialize Python tool: {e}"
            )

        # File reader
        tools.append(
            ReadLocalFileTool()
        )

        # Excel
        tools.append(
            ReadExcelTool()
        )

        # File information
        tools.append(
            FileInfoTool()
        )

        # YouTube transcript
        tools.append(
            YouTubeTranscriptTool()
        )

        # Speech-to-text
        try:

            tools.append(
                SpeechToTextTool()
            )

            print(
                "Speech-to-text tool enabled."
            )

        except Exception as e:

            print(
                f"Speech-to-text unavailable: {e}"
            )

        # ----------------------------------------------------
        # CREATE CODE AGENT
        # ----------------------------------------------------

        self.agent = CodeAgent(
            model=self.model,
            tools=tools,
            max_steps=12,
            verbosity_level=1,
        )

        print(
            f"GAIA Agent ready with {len(tools)} tools."
        )

        print("=" * 60)


    # ========================================================
    # RUN AGENT
    # ========================================================

    def __call__(
        self,
        question: str,
        file_path: str | None = None
    ) -> str:

        print("\n")
        print("=" * 60)
        print("NEW GAIA QUESTION")
        print("=" * 60)

        print(question)

        # ----------------------------------------------------
        # Build task
        # ----------------------------------------------------

        task = f"""
You are an expert autonomous agent solving a GAIA Level 1
benchmark question.

Your goal is to produce the EXACT answer required by the
question.

QUESTION:

{question}

------------------------------------------------------------

IMPORTANT RULES

1. Carefully read the entire question.

2. Determine exactly what the question is asking.

3. If external information is required:
   use web search.

4. If a webpage must be inspected:
   use the webpage tool.

5. If calculations are required:
   use Python.

6. If a spreadsheet is provided:
   inspect it with the Excel tool and Python/pandas.

7. If a Python file is provided:
   read and execute/analyze it with Python.

8. If an audio file is provided:
   use speech-to-text if necessary.

9. If the question contains a YouTube URL:
   try the YouTube transcript tool when the question
   concerns spoken dialogue.

10. Never guess when the information can be obtained
    from a tool.

11. Verify important calculations.

12. Follow the requested output format EXACTLY.

13. Pay attention to:
    - capitalization
    - commas
    - ordering
    - decimal places
    - units
    - first name vs surname
    - city vs country
    - IOC codes
    - algebraic chess notation
    - requested number of words

14. Do not add explanations to the final answer.

15. Do not write:
    "The answer is..."

16. Do not write:
    "FINAL ANSWER"

17. Return ONLY the answer requested by the question.

------------------------------------------------------------
"""

        # ----------------------------------------------------
        # Add file information
        # ----------------------------------------------------

        if file_path:

            file_extension = Path(
                file_path
            ).suffix.lower()

            task += f"""

A FILE IS ATTACHED TO THIS QUESTION.

Local file path:

{file_path}

File extension:

{file_extension}

You MUST inspect the file when it is relevant.

Available file-related tools include:

- file_information
- read_local_file
- read_excel_file
- Python

If the file is XLSX:
use read_excel_file and/or pandas.

If the file is Python:
read the code and execute/analyze it.

If the file is text:
read it.

If the file is audio:
use speech-to-text.

Do not ignore the attached file.
"""

        # ----------------------------------------------------
        # Run
        # ----------------------------------------------------

        try:

            result = self.agent.run(task)

            answer = str(result).strip()

            # ------------------------------------------------
            # Clean accidental formatting
            # ------------------------------------------------

            answer = clean_final_answer(
                answer
            )

            print(
                "FINAL SUBMITTED ANSWER:"
            )

            print(answer)

            print("=" * 60)

            return answer

        except Exception as e:

            print(
                f"Agent execution error: {e}"
            )

            return (
                f"ERROR: {e}"
            )


# ============================================================
# CLEAN FINAL ANSWER
# ============================================================

def clean_final_answer(answer: str) -> str:

    answer = answer.strip()

    # Remove common accidental prefixes

    prefixes = [
        "FINAL ANSWER:",
        "FINAL ANSWER",
        "Answer:",
        "ANSWER:",
        "The answer is:",
        "The answer is"
    ]

    for prefix in prefixes:

        if answer.lower().startswith(
            prefix.lower()
        ):

            answer = answer[
                len(prefix):
            ].strip()

    # Remove markdown code fences

    if answer.startswith("```"):

        lines = answer.splitlines()

        if len(lines) >= 3:

            lines = lines[1:-1]

            answer = "\n".join(
                lines
            ).strip()

    return answer


# ============================================================
# DOWNLOAD GAIA FILE
# ============================================================

def download_task_file(
    task_id: str,
    file_name: str
):

    if not file_name:

        return None

    try:

        file_url = (
            f"{DEFAULT_API_URL}/files/{task_id}"
        )

        print(
            f"Downloading attachment from: {file_url}"
        )

        response = requests.get(
            file_url,
            timeout=60
        )

        response.raise_for_status()

        # Temporary directory
        temp_dir = (
            Path(tempfile.gettempdir())
            / "gaia_files"
        )

        temp_dir.mkdir(
            parents=True,
            exist_ok=True
        )

        safe_name = Path(
            file_name
        ).name

        file_path = (
            temp_dir / safe_name
        )

        file_path.write_bytes(
            response.content
        )

        print(
            f"Downloaded: {file_path}"
        )

        return str(file_path)

    except Exception as e:

        print(
            f"File download failed: {e}"
        )

        return None


# ============================================================
# FETCH QUESTIONS
# ============================================================

def fetch_questions():

    questions_url = (
        f"{DEFAULT_API_URL}/questions"
    )

    print(
        f"Fetching questions from: {questions_url}"
    )

    response = requests.get(
        questions_url,
        timeout=30
    )

    response.raise_for_status()

    data = response.json()

    if not isinstance(data, list):

        raise ValueError(
            "Questions API did not return a list."
        )

    print(
        f"Fetched {len(data)} questions."
    )

    return data


# ============================================================
# RUN EVERYTHING AND SUBMIT
# ============================================================

def run_and_submit_all(
    profile: gr.OAuthProfile | None
):

    # --------------------------------------------------------
    # Check login
    # --------------------------------------------------------

    if profile:

        username = str(
            profile.username
        )

        print(
            f"Logged in user: {username}"
        )

    else:

        print(
            "User is not logged in."
        )

        return (
            "Please login to Hugging Face first.",
            None
        )

    # --------------------------------------------------------
    # Get Space ID
    # --------------------------------------------------------

    space_id = os.getenv(
        "SPACE_ID"
    )

    if not space_id:

        return (
            "SPACE_ID environment variable was not found. "
            "Run this inside your Hugging Face Space.",
            None
        )

    agent_code = (
        f"https://huggingface.co/spaces/"
        f"{space_id}/tree/main"
    )

    print(
        f"Agent code: {agent_code}"
    )

    # --------------------------------------------------------
    # Initialize agent
    # --------------------------------------------------------

    try:

        agent = BasicAgent()

    except Exception as e:

        print(
            f"Agent initialization failed: {e}"
        )

        return (
            f"Agent initialization failed: {e}",
            None
        )

    # --------------------------------------------------------
    # Fetch questions
    # --------------------------------------------------------

    try:

        questions_data = (
            fetch_questions()
        )

    except Exception as e:

        return (
            f"Could not fetch questions: {e}",
            None
        )

    # --------------------------------------------------------
    # Process questions
    # --------------------------------------------------------

    results_log = []

    answers_payload = []

    total_questions = len(
        questions_data
    )

    print(
        f"Processing {total_questions} questions..."
    )

    for index, item in enumerate(
        questions_data,
        start=1
    ):

        print("\n")
        print(
            f"QUESTION {index}/{total_questions}"
        )

        task_id = item.get(
            "task_id"
        )

        question_text = item.get(
            "question"
        )

        file_name = item.get(
            "file_name",
            ""
        )

        # ----------------------------------------------------
        # Validate
        # ----------------------------------------------------

        if not task_id:

            print(
                "Skipping: missing task_id"
            )

            continue

        if question_text is None:

            print(
                "Skipping: missing question"
            )

            continue

        # ----------------------------------------------------
        # Download attachment
        # ----------------------------------------------------

        file_path = None

        if file_name:

            file_path = download_task_file(
                task_id,
                file_name
            )

        # ----------------------------------------------------
        # Run agent
        # ----------------------------------------------------

        try:

            submitted_answer = agent(
                question_text,
                file_path
            )

            # Safety conversion
            submitted_answer = str(
                submitted_answer
            ).strip()

            answers_payload.append(
                {
                    "task_id": task_id,
                    "submitted_answer": submitted_answer
                }
            )

            results_log.append(
                {
                    "Task ID": task_id,
                    "Question": question_text,
                    "Submitted Answer": submitted_answer
                }
            )

            print(
                f"Question {index} completed."
            )

        except Exception as e:

            print(
                f"Question failed: {e}"
            )

            results_log.append(
                {
                    "Task ID": task_id,
                    "Question": question_text,
                    "Submitted Answer":
                        f"AGENT ERROR: {e}"
                }
            )

    # --------------------------------------------------------
    # Check answers
    # --------------------------------------------------------

    if not answers_payload:

        return (
            "Agent produced no answers.",
            pd.DataFrame(
                results_log
            )
        )

    # --------------------------------------------------------
    # Submission
    # --------------------------------------------------------

    submission_data = {
        "username": username.strip(),
        "agent_code": agent_code,
        "answers": answers_payload
    }

    submit_url = (
        f"{DEFAULT_API_URL}/submit"
    )

    print("\n")
    print("=" * 60)
    print(
        f"Submitting {len(answers_payload)} answers"
    )
    print("=" * 60)

    try:

        response = requests.post(
            submit_url,
            json=submission_data,
            timeout=300
        )

        response.raise_for_status()

        result_data = response.json()

        score = result_data.get(
            "score",
            "N/A"
        )

        correct_count = result_data.get(
            "correct_count",
            "?"
        )

        total_attempted = result_data.get(
            "total_attempted",
            "?"
        )

        message = result_data.get(
            "message",
            "No message."
        )

        final_status = (
            "Submission Successful!\n\n"
            f"User: {result_data.get('username', username)}\n"
            f"Score: {score}%\n"
            f"Correct: "
            f"{correct_count}/"
            f"{total_attempted}\n\n"
            f"Message: {message}"
        )

        print(
            final_status
        )

        return (
            final_status,
            pd.DataFrame(
                results_log
            )
        )

    except requests.exceptions.HTTPError as e:

        detail = (
            f"HTTP error {e.response.status_code}"
        )

        try:

            error_json = (
                e.response.json()
            )

            detail += (
                f": {error_json}"
            )

        except Exception:

            detail += (
                f": {e.response.text[:1000]}"
            )

        return (
            f"Submission failed: {detail}",
            pd.DataFrame(
                results_log
            )
        )

    except requests.exceptions.Timeout:

        return (
            "Submission failed: "
            "request timed out. "
            "The agent may have taken too long.",
            pd.DataFrame(
                results_log
            )
        )

    except requests.exceptions.RequestException as e:

        return (
            f"Submission network error: {e}",
            pd.DataFrame(
                results_log
            )
        )

    except Exception as e:

        return (
            f"Unexpected submission error: {e}",
            pd.DataFrame(
                results_log
            )
        )


# ============================================================
# GRADIO INTERFACE
# ============================================================

with gr.Blocks() as demo:

    gr.Markdown(
        "# 🤖 GAIA Agent Evaluation"
    )

    gr.Markdown(
        """
        ### Final Agent Assignment

        This Space runs an autonomous agent against the
        GAIA evaluation questions.

        **Capabilities**

        - 🌐 Web search
        - 🔗 Webpage reading
        - 🐍 Python calculations
        - 📊 Excel analysis
        - 📁 File analysis
        - 🎵 Speech-to-text
        - ▶️ YouTube transcript extraction
        - 🤖 Multi-step reasoning

        **Important:** The final agent response is submitted
        as the answer, so it must follow the exact format
        requested by each question.
        """
    )

    gr.Markdown(
        """
        ### Instructions

        1. Log in to Hugging Face.
        2. Click **Run Evaluation & Submit All Answers**.
        3. The agent will retrieve the GAIA questions.
        4. The agent will solve each question.
        5. Attachments will be downloaded automatically.
        6. Answers will be submitted to the evaluation API.
        7. Your score will appear below.
        """
    )

    gr.Markdown(
        "---"
    )

    gr.LoginButton()

    run_button = gr.Button(
        "🚀 Run Evaluation & Submit All Answers",
        variant="primary"
    )

    status_output = gr.Textbox(
        label="Run Status / Submission Result",
        lines=8,
        interactive=False
    )

    results_table = gr.DataFrame(
        label="Questions and Agent Answers",
        wrap=True
    )

    run_button.click(
        fn=run_and_submit_all,
        outputs=[
            status_output,
            results_table
        ]
    )


# ============================================================
# START APPLICATION
# ============================================================

if __name__ == "__main__":

    print(
        "\n"
        + "-" * 30
        + " GAIA Agent Starting "
        + "-" * 30
    )

    space_host = os.getenv(
        "SPACE_HOST"
    )

    space_id = os.getenv(
        "SPACE_ID"
    )

    if space_host:

        print(
            f"SPACE_HOST: {space_host}"
        )

        print(
            "Runtime URL: "
            f"https://{space_host}.hf.space"
        )

    else:

        print(
            "SPACE_HOST not found."
        )

    if space_id:

        print(
            f"SPACE_ID: {space_id}"
        )

        print(
            "Repository:"
        )

        print(
            f"https://huggingface.co/spaces/{space_id}"
        )

        print(
            "Code:"
        )

        print(
            f"https://huggingface.co/spaces/"
            f"{space_id}/tree/main"
        )

    else:

        print(
            "SPACE_ID not found."
        )

    print(
        "-" * 70
    )

    demo.launch(
        debug=True,
        share=False
    )