import os
import logging
import mimetypes
import subprocess

from typing import Any, List

import gradio as gr
import requests
import pandas as pd
import io
import torchaudio
import torchaudio.transforms as T
import whisper

from llama_index.core.agent.workflow import AgentWorkflow, ToolCallResult, ToolCall, AgentOutput
from llama_index.core.base.llms.types import ChatMessage, TextBlock, ImageBlock, AudioBlock
from llama_index.llms.openai import OpenAI

from agents.video_analyzer_agent import initialize_video_analyzer_agent

os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Assuming agent initializers are in the same directory or a known path
# Adjust import paths if necessary based on deployment structure
try:
    # Existing agents
    from agents.image_analyzer_agent import initialize_image_analyzer_agent
    from agents.reasoning_agent import initialize_reasoning_agent
    from agents.text_analyzer_agent import initialize_text_analyzer_agent
    from agents.code_agent import initialize_code_agent
    from agents.math_agent import initialize_math_agent
    from agents.planner_agent import initialize_planner_agent
    from agents.research_agent import initialize_research_agent
    from agents.role_agent import initialize_role_agent
    # New agents
    from agents.advanced_validation_agent import initialize_advanced_validation_agent
    from agents.long_context_management_agent import initialize_long_context_management_agent
    from agents.synthesis_agent import initialize_synthesis_agent
    AGENT_IMPORT_PATH = "local"
except ImportError as e:
    print(f"Import Error: Could not find agent modules. Tried local and final_project paths. Error: {e}")
    # Set initializers to None or raise error to prevent app start
    initialize_image_analyzer_agent = None
    # ... set all others to None ...
    raise RuntimeError(f"Failed to import agent modules: {e}")

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# --- Constants ---
DEFAULT_API_URL = os.getenv("GAIA_API_URL", "https://agents-course-unit4-scoring.hf.space")

# --- Helper Functions ---
_whisper_model = whisper.load_model("small")

def transcribe_audio(audio_bytes: bytes) -> str:
    logger.info(f"Attempting to transcribe audio file")

    file_like = io.BytesIO(audio_bytes)

    waveform, sample_rate = torchaudio.load(file_like)

    waveform = waveform.mean(dim=0, keepdim=True)  # [1, samples]

    if sample_rate != 16000:
        resampler = T.Resample(orig_freq=sample_rate, new_freq=16000)
        waveform = resampler(waveform)

    waveform = waveform.squeeze(0)

    print(f"Tensor shape : {waveform.shape}, Frequency : {sample_rate} Hz")

    # Load the Whisper model (lazy loading)
    model: whisper.Whisper = _whisper_model  # Uses default size "base" or WHISPER_MODEL_SIZE env var
    if model is None:
        return "Error: Failed to load Whisper model."

    try:
        # Perform transcription
        # The transcribe function handles various audio formats via ffmpeg
        result = whisper.transcribe(model=model, audio=waveform)
        transcribed_text = result["text"]
        detected_language = result.get("language", "unknown")  # Get detected language if available
        logger.info(
            f"Audio transcription successful. Detected language: {detected_language}. Text length: {len(transcribed_text)}")
        return transcribed_text

    except Exception as e:
        # Check if it might be an ffmpeg issue
        if "ffmpeg" in str(e).lower():
            logger.error(f"Error during transcription, possibly ffmpeg issue: {e}", exc_info=True)
            # Check if ffmpeg is installed using shell command
            try:
                subprocess.run(["ffmpeg", "-version"], check=True, capture_output=True)
                # If ffmpeg is installed, the error is likely something else
                return f"Error during transcription (ffmpeg seems installed): {e}"
            except (FileNotFoundError, subprocess.CalledProcessError):
                logger.error("ffmpeg command not found or failed. Please ensure ffmpeg is installed and in PATH.")
                return "Error: ffmpeg not found or not working. Please install ffmpeg."
        else:
            logger.error(f"Unexpected error during transcription: {e}", exc_info=True)
            return f"Error during transcription: {e}"

# --- Agent Initialization (Singleton Pattern) ---
# Initialize the agent workflow once
AGENT_WORKFLOW = None
try:
    logger.info(f"Initializing GAIA Multi-Agent Workflow (import path: {AGENT_IMPORT_PATH})...")
    # Existing agents
    # role_agent = initialize_role_agent()
    code_agent = initialize_code_agent()
    math_agent = initialize_math_agent()
    planner_agent = initialize_planner_agent()
    research_agent = initialize_research_agent()
    text_analyzer_agent = initialize_text_analyzer_agent()
    image_analyzer_agent = initialize_image_analyzer_agent()
    reasoning_agent = initialize_reasoning_agent()
    # New agents
    advanced_validation_agent = initialize_advanced_validation_agent()
    long_context_management_agent = initialize_long_context_management_agent()
    video_analyzer_agent = initialize_video_analyzer_agent()
    synthesis_agent = initialize_synthesis_agent()

    # Check if all agents initialized successfully
    all_agents = [
        code_agent, math_agent, planner_agent, research_agent,
        text_analyzer_agent, image_analyzer_agent, reasoning_agent,
        advanced_validation_agent, long_context_management_agent,
        video_analyzer_agent, synthesis_agent
    ]
    if not all(all_agents):
        raise RuntimeError("One or more agents failed to initialize.")

    AGENT_WORKFLOW = AgentWorkflow(
        agents=all_agents,
        root_agent="reasoning_agent", # Keep planner as root as per plan
        initial_state={
            "research_content": []
        }
    )
    logger.info("GAIA Multi-Agent Workflow initialized successfully.")
except Exception as e:
    logger.error(f"FATAL: Error initializing agent workflow: {e}", exc_info=True)
    # AGENT_WORKFLOW remains None, BasicAgent init will fail

# --- Basic Agent Definition (Wrapper for Workflow) ---
class BasicAgent:
    def __init__(self, workflow: AgentWorkflow):
        if workflow is None:
            logger.error("AgentWorkflow is None, initialization likely failed.")
            raise RuntimeError("AgentWorkflow failed to initialize. Check logs for details.")
        self.agent_workflow = workflow
        logger.info("BasicAgent wrapper initialized.")

    async def __call__(self, question: str | ChatMessage) -> Any:
        if isinstance(question, ChatMessage):
            log_question = str(question.blocks[0].text)[:100] if question.blocks and hasattr(question.blocks[0], "text") else str(question)[:100]
            logger.info(f"Agent received question (first 100 chars): {log_question}...")
        else:
            logger.info(f"Agent received question (first 100 chars): {question[:100]}...")

        handler = self.agent_workflow.run(user_msg=question)

        current_agent = None
        async for event in handler.stream_events():
            if (
                    hasattr(event, "current_agent_name")
                    and event.current_agent_name != current_agent
            ):
                current_agent = event.current_agent_name
                logger.info(f"{'=' * 50}")
                logger.info(f"🤖 Agent: {current_agent}")
                logger.info(f"{'=' * 50}\n")

            # Optional detailed logging (uncomment if needed)
            # from llama_index.core.agent.runner.base import AgentStream, AgentInput
            # if isinstance(event, AgentStream):
            #     if event.delta:
            #         logger.debug(f"STREAM: {event.delta}") # Use debug level
            # elif isinstance(event, AgentInput):
            #     logger.debug(f"📥 Input: {event.input}") # Use debug level
            elif isinstance(event, AgentOutput):
                if event.response and hasattr(event.response, 'content') and event.response.content:
                    logger.info(f"📤 Output: {event.response.content}")
                if event.tool_calls:
                    logger.info(
                        f"🛠️  Planning to use tools: {[call.tool_name for call in event.tool_calls]}"
                    )
            elif isinstance(event, ToolCallResult):
                logger.info(f"🔧 Tool Result ({event.tool_name}):")
                logger.info(f"  Arguments: {event.tool_kwargs}")
                # Limit output logging length if potentially very long
                output_str = str(event.tool_output)
                logger.info(f"  Output: {output_str[:500]}{'...' if len(output_str) > 500 else ''}")
            elif isinstance(event, ToolCall):
                logger.info(f"🔨 Calling Tool: {event.tool_name}")
                logger.info(f"  With arguments: {event.tool_kwargs}")

        answer = await handler
        final_content = answer.response.content if hasattr(answer, 'response') and hasattr(answer.response, 'content') else str(answer)
        logger.info(f"Agent returning final answer: {final_content[:500]}{'...' if len(final_content) > 500 else ''}")
        return answer.response # Return the actual response object expected by Gradio

system_prompt = """
You are a general AI assistant. 
I will give you a result, and with it you will have to transform it to follow the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. 
YOUR FINAL ANSWER should be a number OR 1 or 2 word(s) 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.
If the result is enclosed in double quotes (""), extract and return only what is inside the quotes, applying the formatting rules if needed.

You must never return a full sentence as the final answer. A sentence is strictly forbidden under all circumstances.
"""

llm = OpenAI(model="gpt-4o-mini", api_key=os.getenv("OPENAI_API_KEY"), temperature=0.05, system_prompt=system_prompt)

# --- Helper Functions for run_and_submit_all ---

async def fetch_questions(questions_url: str) -> List[dict] | None:
    """Fetches questions from the GAIA benchmark API."""
    logger.info(f"Fetching questions from: {questions_url}")
    try:
        response = requests.get(questions_url, timeout=30) # Increased timeout
        response.raise_for_status()
        questions_data = response.json()
        if not questions_data:
             logger.warning("Fetched questions list is empty.")
             return None
        logger.info(f"Fetched {len(questions_data)} questions.")
        return questions_data
    except requests.exceptions.RequestException as e:
        logger.error(f"Error fetching questions: {e}", exc_info=True)
        return None
    except requests.exceptions.JSONDecodeError as e:
         logger.error(f"Error decoding JSON response from questions endpoint: {e}", exc_info=True)
         logger.error(f"Response text: {response.text[:500]}")
         return None
    except Exception as e:
        logger.error(f"An unexpected error occurred fetching questions: {e}", exc_info=True)
        return None

async def process_question(agent: BasicAgent, item: dict, base_fetch_file_url: str) -> dict | None:
    """Processes a single question item using the agent."""
    task_id = item.get("task_id")
    question_text = item.get("question")
    file_name = item.get("file_name")

    if not task_id or question_text is None:
        logger.warning(f"Skipping item with missing task_id or question: {item}")
        return None

    message: ChatMessage
    if file_name:
        fetch_file_url = f"{base_fetch_file_url}/{task_id}"
        logger.info(f"Fetching file '{file_name}' for task {task_id} from {fetch_file_url}")
        try:
            response = requests.get(fetch_file_url, timeout=60) # Increased timeout for files
            response.raise_for_status()
            mime_type, _ = mimetypes.guess_type(file_name)
            logger.info(f"File '{file_name}' MIME type guessed as: {mime_type}")

            file_block: TextBlock | ImageBlock | AudioBlock | None = None
            if mime_type:
                # Prioritize specific extensions for text-like content
                text_extensions = (
                    ".txt", ".json", ".xml", ".yaml", ".yml", ".ini", ".cfg", ".toml", ".log", ".properties",
                    ".html", ".htm", ".xhtml", ".css", ".scss", ".sass", ".less", ".svg", ".md", ".rst",
                    ".py", ".js", ".java", ".c", ".cpp", ".h", ".hpp", ".cs", ".go", ".php", ".rb", ".swift", ".kt",
                    ".sh", ".bat", ".ipynb", ".Rmd", ".tex"  # Added more code/markup types
                )
                if mime_type.startswith('text/') or file_name.lower().endswith(text_extensions):
                    try:
                        file_content = response.content.decode('utf-8') # Try UTF-8 first
                    except UnicodeDecodeError:
                        try:
                            file_content = response.content.decode('latin-1') # Fallback
                            logger.warning(f"Decoded file {file_name} using latin-1 fallback.")
                        except Exception as decode_err:
                            logger.error(f"Could not decode file {file_name}: {decode_err}")
                            file_content = f"[Error: Could not decode file content for {file_name}]"
                    file_block = TextBlock(block_type="text", text=f"[File: {file_name}]\n[Content]:\n{file_content}")
                elif mime_type.startswith('image/'):
                    # Pass image content directly for multi-modal models
                    file_block = ImageBlock(url=fetch_file_url, image=response.content)
                elif mime_type.startswith('audio/'):
                    # Pass audio content directly
                    audio_text = transcribe_audio(response.content)
                    file_block = TextBlock(text=f"[Transcribed Audio: {audio_text}]")
                elif mime_type == 'application/pdf':
                    # PDF: Pass a text block indicating the URL for agents to handle
                    logger.info(f"PDF file detected: {file_name}. Passing reference URL.")
                    file_block = TextBlock(text=f"[Reference PDF file available at: {fetch_file_url}]")
                elif file_name.lower().endswith((".xlsx", ".xls", ".csv")):
                    logger.info(f"Data file detected: {file_name}. Passing reference URL.")
                    file_block = TextBlock(text=f"[Reference Data file available at: {fetch_file_url}]")

                # Add handling for other types like video if needed
                # elif mime_type.startswith('video/'):
                #     logger.info(f"Video file detected: {file_name}. Passing reference URL.")
                #     file_block = TextBlock(text=f"[Reference Video file available at: {fetch_file_url}]")

            if file_block:
                 blocks = [TextBlock(text=question_text), file_block]
                 message = ChatMessage(role="user", blocks=blocks)
            else:
                 logger.warning(f"File type for '{file_name}' (MIME: {mime_type}) not directly supported for block creation or no block created (e.g., unsupported). Passing text question only.")
                 message = ChatMessage(role="user", blocks=[TextBlock(text=question_text)])

        except requests.exceptions.RequestException as e:
            logger.error(f"Error fetching file for task {task_id}: {e}", exc_info=True)
            return {"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: Failed to fetch file {file_name} - {e}"}
        except Exception as e:
             logger.error(f"Error processing file for task {task_id}: {e}", exc_info=True)
             return {"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: Failed to process file {file_name} - {e}"}
    else:
        # No file associated with the question
        message = ChatMessage(role="user", blocks=[TextBlock(text=question_text)])

    # Run the agent on the prepared message
    try:
        logger.info(f"Running agent on task {task_id}...")
        submitted_answer_response = await agent(message)
        # Extract content safely
        submitted_answer = submitted_answer_response.content if hasattr(submitted_answer_response, 'content') else str(submitted_answer_response)

        prompt = f"""
        You are a general AI assistant. 
        I will give you a result, and with it you will have to transform it to follow the following template: [YOUR FINAL ANSWER]. 
        YOUR FINAL ANSWER should be a number OR 1 or 2 word(s) 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.
        If the result is enclosed in double quotes (""), extract and return only what is inside the quotes, applying the formatting rules if needed.
        
        You must never return a full sentence as the final answer. A sentence is strictly forbidden under all circumstances.
        QUESTION: {question_text}
        ANSWER: {submitted_answer}
        INSTRUCTIONS: Based on the provided question and answer, generate a final answer that is clear, concise, and directly addresses the question.
        [YOUR FINAL ANSWER]
        """

        final_answer = llm.complete(prompt)

        logger.info(f"👍 Agent submitted answer for task {task_id}: {final_answer.text[:200]}{'...' if len(final_answer.text) > 200 else ''}")
        return {"Task ID": task_id, "Question": question_text, "Submitted Answer": final_answer.text}
    except Exception as e:
         logger.error(f"Error running agent on task {task_id}: {e}", exc_info=True)
         return {"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"}

async def run_and_submit_all( profile: gr.OAuthProfile | None):
    """
    Fetches all questions, runs the BasicAgent on them, submits all answers,
    and displays the results.
    """
    # --- Determine HF Space Runtime URL and Repo URL ---
    space_id = os.getenv("SPACE_ID") # Get the SPACE_ID for sending link to the code

    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"
    fetch_file_url = f"{api_url}/files"

    results_log = []
    answers_payload = []

    try:
        agent = BasicAgent(AGENT_WORKFLOW)
    except Exception as e:
        print(f"Error instantiating agent: {e}")
        return f"Error initializing agent: {e}", None
    # In the case of an app running as a hugging Face space, this link points toward your codebase ( usefull for others so please keep it public)
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
    print(agent_code)

    questions_data = await fetch_questions(questions_url)
    if not questions_data:
        return "Failed to fetch questions.", None

    # 3. Process Questions
    # questions_data = [questions_data[3]]
    for item in questions_data:
        answers = await process_question(agent, item, fetch_file_url)
        results_log.append(answers)
        answers_payload.append({"task_id": answers["Task ID"], "submitted_answer": answers["Submitted Answer"]})

    if not answers_payload:
        print("Agent did not produce any answers to submit.")
        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)

    # 4. Prepare Submission
    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)

    # 5. Submit
    print(f"Submitting {len(answers_payload)} answers to: {submit_url}")

    try:
        response = requests.post(submit_url, json=submission_data, timeout=120) # Increased timeout
        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.')}"
        )
        logger.info("Submission successful.")
        results_df = pd.DataFrame(results_log)
        return final_status, results_df
    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}"
        logger.error(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except requests.exceptions.Timeout:
        status_message = "Submission Failed: The request timed out."
        logger.error(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except requests.exceptions.RequestException as e:
        status_message = f"Submission Failed: Network error - {e}"
        logger.error(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except Exception as e:
        status_message = f"Submission Failed: An unexpected error occurred during submission - {e}"
        logger.error(status_message, exc_info=True)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df

# --- Gradio Interface ---
def create_gradio_interface():
    """Creates and returns the Gradio interface."""
    # --- 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]
        )

    return demo

# --- Main Execution ---
if __name__ == "__main__":
    if not AGENT_WORKFLOW:
        print("ERROR: Agent Workflow failed to initialize. Cannot start Gradio app.")
        print("Please check logs for initialization errors (e.g., missing API keys, import issues).")
    else:
        gradio_app = create_gradio_interface()
        # Launch Gradio app
        # Share=True creates a public link (use with caution)
        # Set server_name="0.0.0.0" to allow access from network
        gradio_app.launch(server_name="0.0.0.0", server_port=7860)