import re
from typing import Union


def extract_final_answer(result: Union[str, dict]) -> str:
    """
    Extract the final answer from the agent's result, removing explanations.
    GAIA requires concise, properly formatted answers.

    Args:
        result: The full result from the agent, either a string or a dictionary

    Returns:
        Extracted final answer
    """
    # Handle dictionary input
    if isinstance(result, dict):
        # Check for final_answer_text first (from agent output)
        if "final_answer_text" in result:
            return str(result["final_answer_text"])
        # Fall back to final_answer key
        if "final_answer" in result:
            return str(result["final_answer"])
        return "No final answer found in result"

    # Handle string input (original logic)
    # First check if there's a specific final_answer marker
    if "final_answer(" in result:
        # Try to extract the answer from final_answer call
        pattern = r"final_answer\(['\"](.*?)['\"]\)"
        matches = re.findall(pattern, result)
        if matches:
            return matches[-1]  # Return the last final_answer if multiple exist

    # If no final_answer marker, look for lines that might contain the answer
    lines = result.strip().split("\n")

    # Check for typical patterns indicating a final answer
    for line in reversed(lines):  # Start from the end
        line = line.strip()

        # Skip empty lines
        if not line:
            continue

        # Look for patterns like "Answer:", "Final answer:", etc.
        if re.match(r"^(answer|final answer|result):?\s+", line.lower()):
            return line.split(":", 1)[1].strip()

        # Check for answers that are comma-separated lists (common in GAIA)
        if (
            "," in line
            and len(line.split(",")) > 1
            and not line.startswith("#")
            and not line.startswith("print(")
        ):
            # It might be a comma-separated list answer
            return line

    # If no clear answer pattern is found, return the last non-empty line
    # (often the answer is simply the last output)
    for line in reversed(lines):
        if (
            line.strip()
            and not line.strip().startswith("#")
            and not line.strip().startswith("print(")
        ):
            return line.strip()

    return "No answer found"