import os
import json
import gradio as gr
import requests
import inspect
import pandas as pd
import asyncio
from typing import ClassVar
from llama_index.core.agent.workflow import FunctionAgent
from pathlib import Path
from llama_index.readers.youtube_transcript import YoutubeTranscriptReader
from llama_index.readers.assemblyai import AssemblyAIAudioTranscriptReader
from llama_index.core import SimpleDirectoryReader
from llama_index.readers.json import JSONReader
from llama_index.readers.pdb import PdbAbstractReader
from llama_index.readers.file import (
    DocxReader,
    HWPReader,
    PDFReader,
    EpubReader,
    FlatReader,
    HTMLTagReader,
    ImageCaptionReader,
    ImageReader,
    ImageVisionLLMReader,
    IPYNBReader,
    MarkdownReader,
    MboxReader,
    PptxReader,
    PandasCSVReader,
    VideoAudioReader,
    UnstructuredReader,
    PyMuPDFReader,
    ImageTabularChartReader,
    XMLReader,
    PagedCSVReader,
    CSVReader,
    RTFReader,
)
from llama_index.core.tools import FunctionTool
from llama_index.llms.openai import OpenAI
from tavily import AsyncTavilyClient
from dotenv import load_dotenv

load_dotenv()

# (Keep Constants as is)
# --- Constants ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"

# --- TOOLS ---

METADATA_FILE_PATH = "metadata_level_1.json"
GAIA_FILES_DIR = "gaia_files"

def get_task_file_content(task_id: str) -> str:
    """
    Reads the content of a file associated with a given task_id.
    The file information is retrieved from metadata_level_1.json.
    Supported file types are PDF, JSON, PDB, TXT, DOCX, CSV, PY, PPTX, and images.
    Returns the file content as a string or an error/info message if the file cannot be processed.
    """
    base_dir = Path(__file__).resolve().parent
    gaia_files_full_dir = base_dir / GAIA_FILES_DIR

    if not gaia_files_full_dir.exists() or not gaia_files_full_dir.is_dir():
        return f"Error: GAIA files directory '{gaia_files_full_dir}' not found."

    found_files = list(gaia_files_full_dir.glob(f"{task_id}.*"))

    if not found_files:
        return f"Info: No file found in '{gaia_files_full_dir}' starting with task_id '{task_id}'."
    
    # If multiple files match the pattern (e.g., task_id.txt, task_id.pdf), pick the first one.
    # You might want to add more sophisticated logic here if needed (e.g., based on preferred extension).
    file_path = found_files[0]
    file_name = file_path.name

    if not file_path.is_file(): # Should not happen with glob if directory is not named like a file
        return f"Error: Path '{file_path}' found for task_id '{task_id}' is not a file."

    # Construct the file path relative to the app.py directory
    gaia_files_full_dir = base_dir / GAIA_FILES_DIR
    file_path = gaia_files_full_dir / file_name

    if not file_path.exists():
        return f"Error: File '{file_path}' (for task_id '{task_id}') not found."

    _, extension = os.path.splitext(file_name.lower())
    docs = []
    content = ""

    try:
        if extension == ".json":
            loader = JSONReader()
            docs = loader.load_data(input_file=file_path)
        elif extension == ".pdb": 
            loader = PdbAbstractReader(input_files=[str(file_path)])
            docs = loader.load_data()
        elif extension in [".pdf", ".txt", ".docx", ".csv", ".py"]:
            file_extractor = {
                ".pdf": PDFReader(),
                ".txt": FlatReader(),
                ".docx": DocxReader(),
                ".csv": CSVReader(),
                ".py": FlatReader(),
                # ".pptx": PptxReader()
            }
            loader = SimpleDirectoryReader(input_files=[str(file_path)], file_extractor=file_extractor)
            docs = loader.load_data()
        elif extension in [".png", ".jpg", ".jpeg"]: # Common image types handled by SimpleDirectoryReader
            parser = ImageReader()
            file_extractor = {
                ".jpg": parser,
                ".jpeg": parser,
                ".png": parser,
            }
            loader = SimpleDirectoryReader(input_files=[str(file_path)], file_extractor=file_extractor)
            docs = loader.load_data()
        elif extension in [".mp3", ".wav"]:
            loader = AssemblyAIAudioTranscriptReader(file_path=str(file_path), api_key=os.getenv("ASSEMBLYAI_API_KEY"))
            docs = loader.load_data()
        else:
            return f"Unsupported file type: {extension} for file {file_name}."

        if docs:
            content = ""
            for doc in docs:
                # se tiver get_text(), usa — caso contrário, pega doc.text
                if hasattr(doc, "get_text"):
                    content = "\n".join([content, doc.get_text()])
                else:
                    # .text é o campo padrão de qualquer Document em LlamaIndex
                    content = "\n".join([content, doc.text or ""])
            print(content)
            return content
        return f"Info: No content extracted from file {file_name}. The file might be empty or the loader did not process it."

    except Exception as e:
        return f"Error reading file {file_name} for task_id '{task_id}': {str(e)}"


def get_youtube_transcript(url: str) -> str:
    """
    Given a YouTube URL, fetches its transcript.
    Returns:
        transcript: A string of the transcript text.
    """

    try:
        loader = YoutubeTranscriptReader()
        documents = loader.load_data(
            ytlinks=[url]
        )
        transcript = "\n".join([doc.text for doc in documents])
        return transcript
    except Exception as e:
        return f"Error fetching transcript for URL {url}: {str(e)}"

async def search_web(query: str) -> str:
    """
    Useful for using the web to answer questions.
    Use wisely and do not use more than 3 times per question.

    Args:
        query: The query to search for.
    
    Returns:
        search_results: A string of the search results.
    """
    client = AsyncTavilyClient(api_key=os.environ.get("TAVILY_API_KEY"))
    return str(await client.search(query))

# --- Utils ---
def add_final_answer(row):
    task_id_to_final_answer = {}
    with open(METADATA_FILE_PATH, "r") as f:
        for line in f:
            metadata = json.loads(line)
            task_id_to_final_answer[metadata['task_id']] = metadata['Final answer']

    row['final_answer'] = task_id_to_final_answer.get(row['Task ID'], "")
    return row

# --- Basic Agent Definition ---
# ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
class BasicAgent(FunctionAgent):
    search_web_count: ClassVar[int] = 0
    max_search_web_calls: ClassVar[int] = 3
    def __init__(self):
        BasicAgent.search_web_count = 0
        BasicAgent.max_search_web_calls = 3

        get_task_data_tool = FunctionTool.from_defaults(
            fn=get_task_file_content,
            name="get_task_file_content",
            description=(
                "Reads and returns the content of a file associated with a given task_id. "
                "Use this tool if the question implies needing information from a specific file related to the task. "
                "You must provide the 'task_id' to this tool. The task_id will be part of the input question."
            )
        )
        get_youtube_transcript_tool = FunctionTool.from_defaults(
            fn=get_youtube_transcript,
            name="get_youtube_transcript",
            description=(
                "Reads and returns the transcript of a YouTube video associated with a given URL. "
                "Use this tool if the question implies needing information from a specific YouTube video related to the task. "
                "You must provide the 'url' to this tool. The url will be part of the input question."
            )
        )
        search_web_tool = FunctionTool.from_defaults(
            fn=self._search_web_wrapper,
            name="search_web",
            description=(
                "Useful for using the web to answer questions. "
                "Use wisely and do not use more than 3 times per question."
            )
        )

        super().__init__(
            tools=[get_task_data_tool, get_youtube_transcript_tool, search_web_tool],
            llm=OpenAI(model=os.getenv("OPENAI_MODEL")),
            system_prompt="""
            You are a general AI assistant. I will ask you a question. 
            The question will be prefixed with 'Task ID: <id>. Question: '. 
            If the question requires information from a file associated with this Task ID, 
            extract the <id> and use the 'get_task_file_content' tool, providing the extracted Task ID to it.
            If you can't retrieve the file content or it doesn't exists, try your best to answer the question.

            Your 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.
            """,
            verbose=True
        )
        print("BasicAgent initialized.")

    async def _search_web_wrapper(self, query: str) -> str:
        """
        Wrapper for the search_web tool to limit its usage per question.
        """
        if BasicAgent.search_web_count >= BasicAgent.max_search_web_calls:
            return "Search limit reached. You have already used the web search tool 3 times for this question."
        BasicAgent.search_web_count += 1
        # Call the original/global search_web function
        return await search_web(query)

    def reset_search_web_count(self):
        """Resets the search_web call counter."""
        BasicAgent.search_web_count = 0

    async def __call__(self, question: str) -> str:
        print(f"Agent received question {question}...")
        answer = await self.run(question)
        print(f"Agent returned answer: {answer}")
        return answer

def run_random_one(profile: gr.OAuthProfile | None):
    """
    Fetches a random question, runs the BasicAgent on it, submits the answer,
    and displays the result.
    """
    # --- Determine HF Space Runtime URL and Repo URL ---
    api_url = os.getenv("API_URL") or DEFAULT_API_URL

    if profile:
        username= f"{profile.username}"
        print(f"User logged in: {username}")
    else:
        print("User not logged in.")
        return "Please Login to Hugging Face with the button.", None

    # --- Fetch Random Question ---
    random_question_url = f"{api_url}/random-question"
    response = requests.get(random_question_url)
    if response.status_code != 200:
        return "Failed to fetch a random question.", None
    response_data = response.json()
    question = response_data["question"]
    task_id = response_data["task_id"]
    print("----------------")
    print(response_data)
    print("----------------")
    print(task_id)
    print(question)
    print("----------------")

    basic_agent = BasicAgent()
    # Augment the question with task_id context for the agent's tool
    question_for_agent = f"Task ID: {task_id}. Question: {question}"
    print(f"Augmented question for agent: {question_for_agent}") # For debugging
    answer = asyncio.run(basic_agent(question_for_agent))
    answer = answer.response.blocks[0].text
    print("----------------")
    print(answer)
    print("----------------")

    # --- Submit Answer ---
    # submit_url = f"{api_url}/submit"
    # submission_data = {
    #     "username": os.getenv("USERNAME"),
    #     "agent_code": os.getenv("AGENT_CODE"),
    #     "answers": [
    #         {"task_id": task_id, "question": question, "answer": answer}
    #     ]
    # }
    # response = requests.post(submit_url, json=submission_data)
    # if response.status_code != 200:
    #     return "Failed to submit the answer.", None
    # result = response.json()

    # --- Display Result ---
    # print("-------------------------------")
    # print(result)
    # print("-------------------------------")
    return "Success", pd.DataFrame([{"task_id": task_id, "question": question, "answer": answer}])


def _fetch_questions(questions_url: str):
    """Fetches questions from the specified URL."""
    print(f"Fetching questions from: {questions_url}")
    try:
        response = requests.get(questions_url, timeout=15)
        response.raise_for_status()
        questions_data = response.json()
        if not questions_data:
            print("Fetched questions list is empty.")
            return None, "Fetched questions list is empty or invalid format."
        print(f"Fetched {len(questions_data)} questions.")
        return questions_data, None
    except requests.exceptions.RequestException as e:
        print(f"Error fetching questions: {e}")
        return None, f"Error fetching questions: {e}"
    except requests.exceptions.JSONDecodeError as e:
        print(f"Error decoding JSON response from questions endpoint: {e}")
        # It's good practice to check if response exists before accessing .text
        response_text = response.text[:500] if hasattr(response, 'text') else "No response text available"
        print(f"Response text: {response_text}")
        return None, f"Error decoding server response for questions: {e}"
    except Exception as e:
        print(f"An unexpected error occurred fetching questions: {e}")
        return None, f"An unexpected error occurred fetching questions: {e}"

def _run_agent_on_questions(basic_agent: BasicAgent, questions_data: list):
    """Runs the agent on the provided questions and logs results."""
    results_log = []
    answers_payload = []
    print(f"Running agent on {len(questions_data)} questions...")
    for item in questions_data:
        basic_agent.reset_search_web_count()  # Reset counter for each new question
        task_id = item.get("task_id")
        question = item.get("question")
        if not task_id or question is None:
            print(f"Skipping item with missing task_id or question: {item}")
            continue
        try:
            question_for_agent = f"Task ID: {task_id}. Question: {question}"
            # Ensure basic_agent is an async function or handle appropriately
            # For this refactor, assuming basic_agent call is correct as is.
            submitted_answer_obj = asyncio.run(basic_agent(question_for_agent))
            submitted_answer = submitted_answer_obj.response.blocks[0].text
            answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
            results_log.append({"Task ID": task_id, "Question": question, "Submitted Answer": submitted_answer})
        except Exception as e:
            print(f"Error running agent on task {task_id}: {e}")
            results_log.append({"Task ID": task_id, "Question": question, "Submitted Answer": f"AGENT ERROR: {e}"})
    
    if not answers_payload:
        print("Agent did not produce any answers to submit.")
        return None, results_log, "Agent did not produce any answers to submit."
        
    return answers_payload, results_log, None

def _prepare_submission_payload(username: str, agent_code: str, answers_payload: list):
    """Prepares the submission payload."""
    submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
    status_update = f"Agent finished. Submitting {len(answers_payload)} answers for user '{username}'..."
    print(status_update)
    return submission_data

def _submit_and_process_results(submit_url: str, submission_data: dict, results_log: list):
    """Submits answers and processes the results from the server."""
    print(f"Submitting {len(submission_data.get('answers', []))} answers to: {submit_url}")
    try:
        response = requests.post(submit_url, json=submission_data, timeout=60)
        response.raise_for_status()
        result_data = response.json()
        final_status = (
            f"Submission Successful!\n"
            f"User: {result_data.get('username')}\n"
            f"Overall Score: {result_data.get('score', 'N/A')}% "
            f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
            f"Message: {result_data.get('message', 'No message received.')}"
        )
        print("Submission successful.")
        return final_status, results_log
    except requests.exceptions.HTTPError as e:
        error_detail = f"Server responded with status {e.response.status_code}."
        try:
            error_json = e.response.json()
            error_detail += f" Detail: {error_json.get('detail', e.response.text)}"
        except requests.exceptions.JSONDecodeError:
            error_detail += f" Response: {e.response.text[:500]}"
        status_message = f"Submission Failed: {error_detail}"
        print(status_message)
        return status_message, results_log
    except requests.exceptions.Timeout:
        status_message = "Submission Failed: The request timed out."
        print(status_message)
        return status_message, results_log
    except requests.exceptions.RequestException as e:
        status_message = f"Submission Failed: Network error - {e}"
        print(status_message)
        return status_message, results_log
    except Exception as e:
        status_message = f"An unexpected error occurred during submission: {e}"
        print(status_message)
        return status_message, results_log


def run_and_submit_all(profile: gr.OAuthProfile | None):
    """
    Fetches all questions, runs the BasicAgent on them, submits all answers,
    and displays the results.
    """
    space_id = os.getenv("SPACE_ID")

    if profile:
        username = f"{profile.username}"
        print(f"User logged in: {username}")
    else:
        print("User not logged in.")
        return "Please Login to Hugging Face with the button.", None

    api_url = DEFAULT_API_URL
    questions_url = f"{api_url}/questions"
    submit_url = f"{api_url}/submit"

    # 1. Instantiate Agent
    try:
        basic_agent = BasicAgent()
    except Exception as e:
        print(f"Error instantiating agent: {e}")
        return f"Error initializing agent: {e}", None
    
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
    print(f"Agent code URL: {agent_code}")

    # 2. Fetch Questions
    questions_data, error_message = _fetch_questions(questions_url)
    if error_message:
        return error_message, None # No DataFrame to return here

    # 3. Run your Agent
    answers_payload, results_log, error_message = _run_agent_on_questions(basic_agent, questions_data)
    if error_message:
        # If agent produced no answers, results_log might still be useful
        return error_message, pd.DataFrame(results_log if results_log else [])

    # 4. Save logs
    results_log = pd.DataFrame(results_log)
    results_log = results_log.apply(add_final_answer, axis=1)
    results_log.to_csv("results_log.csv", index=False)

    # 5. Prepare Submission
    submission_data = _prepare_submission_payload(username, agent_code, answers_payload)

    # 6. Submit and Process Results
    final_status, results_df = _submit_and_process_results(submit_url, submission_data, results_log)
    return final_status, results_df


# --- Build Gradio Interface using Blocks ---
with gr.Blocks() as demo:
    gr.Markdown("# Basic Agent Evaluation Runner")
    gr.Markdown(
        """
        **Instructions:**

        1.  Please clone this space, then modify the code to define your agent's logic, the tools, the necessary packages, etc ...
        2.  Log in to your Hugging Face account using the button below. This uses your HF username for submission.
        3.  Click 'Run Evaluation & Submit All Answers' to fetch questions, run your agent, submit answers, and see the score.

        ---
        **Disclaimers:**
        Once clicking on the "submit button, it can take quite some time ( this is the time for the agent to go through all the questions).
        This space provides a basic setup and is intentionally sub-optimal to encourage you to develop your own, more robust solution. For instance for the delay process of the submit button, a solution could be to cache the answers and submit in a seperate action or even to answer the questions in async.
        """
    )

    gr.LoginButton()

    run_button = gr.Button("Run Evaluation & Submit All Answers")

    status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
    # Removed max_rows=10 from DataFrame constructor
    results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)


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

    random_button = gr.Button("Run Random Question")
    random_button.click(
        fn=run_random_one,
        outputs=[status_output, results_table]
    )

    # task_id = gr.Textbox(label="Task ID")
    # task_content_output = gr.Textbox(label="Task Content", lines=5, interactive=True)
    # get_task_content_button = gr.Button("Get Task Content")
    # get_task_content_button.click(
    #     fn=get_task_file_content,
    #     inputs=[task_id],
    #     outputs=[task_content_output]
    # )

    # youtube_url = gr.Textbox(label="Youtube URL")
    # youtube_transcript_output = gr.Textbox(label="Youtube Transcript", lines=5, interactive=True)
    # get_youtube_transcript_button = gr.Button("Get Youtube Transcript")
    # get_youtube_transcript_button.click(
    #     fn=get_youtube_transcript,
    #     inputs=[youtube_url],
    #     outputs=[youtube_transcript_output]
    # )

if __name__ == "__main__":
    print("\n" + "-"*30 + " App Starting " + "-"*30)
    # Check for SPACE_HOST and SPACE_ID at startup for information
    space_host_startup = os.getenv("SPACE_HOST")
    space_id_startup = os.getenv("SPACE_ID") # Get SPACE_ID at startup

    if space_host_startup:
        print(f"✅ SPACE_HOST found: {space_host_startup}")
        print(f"   Runtime URL should be: https://{space_host_startup}.hf.space")
    else:
        print("ℹ️  SPACE_HOST environment variable not found (running locally?).")

    if space_id_startup: # Print repo URLs if SPACE_ID is found
        print(f"✅ SPACE_ID found: {space_id_startup}")
        print(f"   Repo URL: https://huggingface.co/spaces/{space_id_startup}")
        print(f"   Repo Tree URL: https://huggingface.co/spaces/{space_id_startup}/tree/main")
    else:
        print("ℹ️  SPACE_ID environment variable not found (running locally?). Repo URL cannot be determined.")

    print("-"*(60 + len(" App Starting ")) + "\n")

    print("Launching Gradio Interface for Agent Evaluation...")
    demo.launch(debug=True, share=False)