import os
import re
import io
import json
import time
import math
import traceback
import contextlib
from typing import TypedDict, Annotated
import operator

import requests
import pandas as pd
import gradio as gr

try:
    import spaces

    @spaces.GPU
    def _zerogpu_warmup():
        # Dummy function so this ZeroGPU Space passes its startup check.
        # This app doesn't need GPU compute (inference runs via Gemini/Groq
        # APIs), so this function is never called for real work.
        return True
except ImportError:
    pass

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.tools import tool
from langchain_google_genai import ChatGoogleGenerativeAI
from groq import Groq as GroqClient
from langchain_community.tools import DuckDuckGoSearchRun, WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper

# =========================================================
# CONSTANTS
# =========================================================
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
MODEL_NAME = "gemini-2.5-flash"
RECURSION_LIMIT = 18          # ceiling on ReAct agent<->tool round-trips per pass
MAX_REFLECTION_ROUNDS = 2     # extra research passes allowed if not yet confident

# =========================================================
# MODELS / CLIENTS
# =========================================================
# Gemini 2.5 Flash handles ALL reasoning, tool-calling, vision, and cleanup.
# Requires GOOGLE_API_KEY as a Space secret (aistudio.google.com/apikey).
llm = ChatGoogleGenerativeAI(model=MODEL_NAME, temperature=0)

# Groq is used ONLY for its free Whisper transcription API.
# Requires GROQ_API_KEY as a Space secret (console.groq.com).
_groq_client = GroqClient(api_key=os.getenv("GROQ_API_KEY"))

# =========================================================
# FORMAT / CLEANUP (GAIA grades by exact string match)
# =========================================================
FORMAT_RULES = """
Formatting rules for the final answer (graded by EXACT STRING MATCH):
- No explanations, no "the answer is", no "FINAL ANSWER:" prefix — just the answer itself.
- If asked for a number, write only the number (no commas, no units, no $ sign) unless the
  question explicitly asks for units.
- If asked for a string, use as few words as possible, no articles ("a", "the") unless the
  question requires them, and do not abbreviate unless asked.
- If asked for a comma separated list, apply the above rules to each element and follow any
  ordering/alphabetization instructions exactly.
- Match the exact capitalization and spelling implied by the question when naming entities.
"""

CLEANUP_PROMPT = (
    "Extract ONLY the final answer from the text below. Strip ALL extra words, units, labels, "
    "explanations, and surrounding context — return the bare answer only. If it's a number, "
    "return digits only (no commas, no $ signs) unless units were explicitly requested. If it's "
    "a name, return just the name. No punctuation unless it is part of the answer itself. "
    "No 'FINAL ANSWER:' prefix. No surrounding quotes."
)

_STRIP_PATTERNS = [
    r'^\s*final answer\s*[:\-]\s*',
    r'^\s*the answer is\s*[:\-]?\s*',
    r'^\s*answer\s*[:\-]\s*',
]


def clean_answer(raw_text: str) -> str:
    """Second-pass LLM call + regex safety net to strip a draft answer down to
    the bare exact-match string GAIA expects."""
    text = str(raw_text).strip()
    try:
        cleaned = llm.invoke([
            SystemMessage(content=CLEANUP_PROMPT),
            HumanMessage(content=text),
        ])
        text = str(cleaned.content).strip()
    except Exception as e:
        print(f"clean_answer LLM pass failed, using regex fallback only: {e}")

    # Regex safety net in case the cleanup call itself left boilerplate in place.
    for pattern in _STRIP_PATTERNS:
        text = re.sub(pattern, '', text, flags=re.IGNORECASE)
    text = text.strip().strip('"').strip("'").strip()
    if text.endswith('.') and not re.search(r'\d\.\d$', text):
        text = text.rstrip('.')
    return text.strip()


# =========================================================
# RAW IMPLEMENTATIONS (wrapped as @tool below)
# =========================================================

def _read_file_impl(file_path: str) -> str:
    if not os.path.exists(file_path):
        return f"Error: file not found at {file_path}"
    ext = file_path.lower().rsplit(".", 1)[-1]
    try:
        if ext == "csv":
            df = pd.read_csv(file_path)
            return (
                f"CSV loaded. Shape: {df.shape}. Columns: {list(df.columns)}\n\n"
                f"Preview:\n{df.head(20).to_string()}\n\n"
                "If the question needs a max/min/sum/average/count/filter/sort, "
                f"use python_tool with pd.read_csv(r'{file_path}') instead of this preview."
            )
        elif ext in ("xlsx", "xls"):
            df = pd.read_excel(file_path)
            return (
                f"Excel loaded. Shape: {df.shape}. Columns: {list(df.columns)}\n\n"
                f"Preview:\n{df.head(20).to_string()}\n\n"
                "If the question needs a max/min/sum/average/count/filter/sort, "
                f"use python_tool with pd.read_excel(r'{file_path}') instead of this preview."
            )
        elif ext == "json":
            with open(file_path, "r") as f:
                data = json.load(f)
            return json.dumps(data, indent=2)[:6000]
        elif ext == "pdf":
            try:
                from pypdf import PdfReader
                reader = PdfReader(file_path)
                text = "\n".join(page.extract_text() or "" for page in reader.pages)
                return text[:10000] if text.strip() else "No extractable text found in PDF (it may be scanned/image-based)."
            except Exception as e:
                return f"Error reading PDF: {e}"
        elif ext in ("txt", "md"):
            with open(file_path, "r", errors="ignore") as f:
                return f.read()[:10000]
        else:
            with open(file_path, "r", errors="ignore") as f:
                return f.read()[:10000]
    except Exception as e:
        return f"Error reading file: {e}"


def _transcribe_audio_impl(file_path: str) -> str:
    if not os.path.exists(file_path):
        return f"Error: file not found at {file_path}"
    try:
        with open(file_path, "rb") as f:
            transcription = _groq_client.audio.transcriptions.create(
                file=(os.path.basename(file_path), f.read()),
                model="whisper-large-v3-turbo",
                response_format="text",
            )
        return str(transcription)
    except Exception as e:
        return f"Error transcribing audio: {e}"


def _analyze_image_impl(file_path: str, question: str) -> str:
    import base64 as b64
    if not os.path.exists(file_path):
        return f"Error: file not found at {file_path}"
    try:
        with open(file_path, "rb") as f:
            img_b64 = b64.b64encode(f.read()).decode("utf-8")
        ext = file_path.lower().rsplit(".", 1)[-1]
        mime = "image/png" if ext == "png" else "image/jpeg"
        response = llm.invoke([
            HumanMessage(content=[
                {"type": "text", "text": question},
                {"type": "image_url", "image_url": f"data:{mime};base64,{img_b64}"},
            ])
        ])
        return str(response.content)
    except Exception as e:
        return f"Error analyzing image: {e}"


def _get_youtube_transcript_impl(url: str) -> str:
    try:
        from youtube_transcript_api import YouTubeTranscriptApi
        match = re.search(r'(?:v=|youtu\.be/)([\w-]+)', url)
        if not match:
            return "Error: could not extract a video ID from that URL."
        video_id = match.group(1)
        transcript = YouTubeTranscriptApi.get_transcript(video_id)
        return " ".join(seg["text"] for seg in transcript)
    except Exception as e:
        return (
            f"Error fetching transcript ({e}). Transcript unavailable — fall back to "
            "web_search for the video's title, description, or discussions of its content."
        )


# =========================================================
# TOOLS — the agent chooses which of these to call, and when.
# =========================================================

@tool
def calculator(expression: str) -> str:
    """Evaluate ONE simple arithmetic expression, e.g. "12 * (3 + 4)" or "156/4 - 2".
    Use this for quick single-line arithmetic. For anything involving dates,
    multi-step logic, data files, or statistics, use python_tool instead."""
    try:
        allowed_chars = set("0123456789+-*/(). %")
        if not all(c in allowed_chars or c.isspace() for c in expression):
            return "Error: expression contains characters this calculator doesn't support. Use python_tool instead."
        return str(eval(expression, {"__builtins__": {}}))
    except Exception as e:
        return f"Error evaluating expression: {e}"


@tool
def python_tool(code: str) -> str:
    """Execute Python code — the preferred way to do ANY calculation you want
    verified rather than done mentally: percentages, statistics, averages,
    dates/time, currency conversion, geometry, counting, or CSV/Excel analysis.
    Pandas (pd), math, and re are pre-imported. If a data file was downloaded,
    load it yourself with pd.read_csv(path) or pd.read_excel(path) and use
    pandas operations (.sum(), .mean(), .max(), .sort_values(), filtering, etc.)
    rather than reasoning over a printed preview. Print anything you want to
    see with print() — only printed output is returned to you."""
    import datetime
    safe_globals = {
        "__builtins__": __builtins__,
        "pd": pd,
        "math": math,
        "re": re,
        "datetime": datetime,
    }
    buffer = io.StringIO()
    try:
        with contextlib.redirect_stdout(buffer):
            exec(code, safe_globals)
        output = buffer.getvalue().strip()
        return output if output else "Code ran with no printed output — use print() to surface results."
    except Exception as e:
        return f"Error executing code: {e}"


_ddg_search = DuckDuckGoSearchRun(name="ddg_search")


@tool
def web_search(query: str) -> str:
    """Search the web via DuckDuckGo for current facts, names, dates, or events
    you aren't fully certain of. If results look thin, empty, or irrelevant,
    call this again with a rewritten query — different keywords, more or less
    specific, or a different angle — rather than giving up after one try."""
    for attempt in range(2):
        try:
            result = _ddg_search.run(query)
            if result and len(result.strip()) > 20:
                return result
        except Exception as e:
            print(f"web_search attempt {attempt + 1} failed: {e}")
            time.sleep(1)
    return (
        "No usable results for this query. Rewrite it with different or more "
        "specific keywords and try again, or try wikipedia_search."
    )


_wiki = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper(top_k_results=2, doc_content_chars_max=3000))


@tool
def wikipedia_search(query: str) -> str:
    """Look up a topic on Wikipedia. Best for encyclopedic facts about people,
    places, organizations, historical events, and concepts."""
    try:
        result = _wiki.run(query)
        return result if result else "No Wikipedia page found for this query. Try web_search instead."
    except Exception as e:
        return f"Wikipedia lookup failed: {e}. Try web_search instead."


@tool
def read_file(file_path: str) -> str:
    """Read a downloaded file's contents — supports csv, xlsx, xls, json, pdf,
    txt, and md. For CSV/Excel this returns a schema + preview only — if the
    question needs a max/min/sum/average/count/filter/sort, use python_tool
    with pandas on this same file_path instead of reasoning over the preview."""
    return _read_file_impl(file_path)


IMAGE_ANALYSIS_HINT = (
    "Examine the image closely and account for whatever is relevant: text "
    "(read it via OCR), object/species/logo identification, exact counts of "
    "items or people, colors, chart or graph values and axis labels, map "
    "locations or routes, table rows and columns, screenshots (UI text, "
    "filenames, timestamps), and small/fine details."
)


@tool
def analyze_image(file_path: str, question: str) -> str:
    """Analyze a downloaded image with vision to answer a specific question
    about it. Pass the exact file_path and a focused question describing
    exactly what to look for (e.g. "What number is on the scoreboard?")."""
    return _analyze_image_impl(file_path, f"{question}\n\n{IMAGE_ANALYSIS_HINT}")


@tool
def transcribe_audio(file_path: str) -> str:
    """Transcribe a downloaded audio file (mp3, wav, m4a, ogg, flac) to text
    via Whisper. This returns the raw transcript only — reason over it
    yourself in a follow-up step rather than treating it as the final answer."""
    return _transcribe_audio_impl(file_path)


@tool
def get_youtube_transcript(url: str) -> str:
    """Fetch the transcript of a YouTube video given its URL. If the
    transcript is unavailable, this returns an error — fall back to
    web_search for information about the video instead."""
    return _get_youtube_transcript_impl(url)


TOOLS = [
    web_search,
    wikipedia_search,
    calculator,
    python_tool,
    read_file,
    analyze_image,
    transcribe_audio,
    get_youtube_transcript,
]
llm_with_tools = llm.bind_tools(TOOLS)

# =========================================================
# SYSTEM PROMPT — one prompt for every question, of every kind.
# There is no separate prompt per question "type": the agent is told what
# tools exist and decides for itself whether and which to use.
# =========================================================
## WITH:
RESEARCH_SYSTEM_PROMPT = f"""You are answering GAIA benchmark questions. You MUST use tools to find answers - NEVER guess.

Available tools: web_search, wikipedia_search, calculator, python_tool, read_file, analyze_image, transcribe_audio, get_youtube_transcript.

CRITICAL RULES:
1. ALWAYS use tools. Never answer from memory or guess.
2. If a file is provided, use read_file/analyze_image/transcribe_audio FIRST.
3. Use python_tool for ALL calculations - never calculate mentally.
4. Use web_search for ANY fact you're not 100% certain about.
5. If a tool fails, try a different approach or tool.
6. Complete your research in 10 steps or less.
7. Your final message must be ONLY the answer - no other text.

{FORMAT_RULES}
"""

# =========================================================
# LANGGRAPH — true ReAct loop (Agent <-> Tools), no upfront routing.
# =========================================================
class AgentState(TypedDict):
    messages: Annotated[list, operator.add]


def call_model(state: AgentState):
    response = llm_with_tools.invoke(state["messages"])
    return {"messages": [response]}


def should_continue(state: AgentState):
    last_message = state["messages"][-1]
    if getattr(last_message, "tool_calls", None):
        return "tools"
    return END


tool_node = ToolNode(TOOLS)
research_graph = StateGraph(AgentState)
research_graph.add_node("agent", call_model)
research_graph.add_node("tools", tool_node)
research_graph.set_entry_point("agent")
research_graph.add_conditional_edges("agent", should_continue, {"tools": "tools", END: END})
research_graph.add_edge("tools", "agent")
compiled_research_graph = research_graph.compile()

def run_react_agent(question: str, context_note: str = "") -> str:
    """Run the ReAct loop once without complex reflection to avoid loops."""
    full_question = f"{context_note}\n\n{question}".strip() if context_note else question
    
    # Force tool usage by making it explicit
    forced_prompt = f"""IMPORTANT: You MUST use your tools to answer this question. Do not guess.{full_question}"""
    
    messages = [SystemMessage(content=RESEARCH_SYSTEM_PROMPT), HumanMessage(content=forced_prompt)]
    
    try:
        result = compiled_research_graph.invoke(
            {"messages": messages}, 
            config={"recursion_limit": RECURSION_LIMIT}
        )
        candidate = str(result["messages"][-1].content).strip()
        
        # If the answer looks like it contains reasoning, try to extract just the final answer
        if len(candidate) > 500:
            # Too long - ask LLM to extract just the answer
            try:
                cleaned = llm.invoke([
                    SystemMessage(content="Extract only the final answer from this text. Return just the answer, nothing else."),
                    HumanMessage(content=candidate)
                ])
                candidate = str(cleaned.content).strip()
            except:
                pass
        
        return clean_answer(candidate) if candidate else "unknown"
        
    except Exception as e:
        print(f"Agent failed: {e}")
        return "unknown"
# =========================================================
# AGENT WRAPPER — same public interface as before (question, task_id) -> answer.
# No classifier: file/video info is passed as CONTEXT, and the LLM decides
# which tool(s), if any, to call.
# =========================================================
class BasicAgent:
    def __init__(self, api_url: str = DEFAULT_API_URL):
        self.api_url = api_url
        print("BasicAgent initialized (single ReAct + reflection pipeline).")

    def _download_file_if_any(self, task_id: str) -> str | None:
        url = f"{self.api_url}/files/{task_id}"
        try:
            resp = requests.get(url, timeout=15)
            if resp.status_code != 200 or not resp.content:
                return None
            cd = resp.headers.get("content-disposition", "")
            match = re.search(r'filename="?([^";]+)"?', cd)
            filename = match.group(1) if match else f"{task_id}_file"
            local_path = os.path.join("/tmp", filename)
            with open(local_path, "wb") as f:
                f.write(resp.content)
            return local_path
        except Exception as e:
            print(f"No file downloaded for {task_id}: {e}")
            return None

    def __call__(self, question: str, task_id: str | None = None) -> str:
        print(f"Processing task {task_id}: {question[:80]}...")

        local_path = self._download_file_if_any(task_id) if task_id else None
        youtube_match = re.search(r'(?:youtube\.com/watch\?v=|youtu\.be/)[\w-]+', question)

        context_lines = []
        if local_path:
            ext = local_path.lower().rsplit(".", 1)[-1]
            if ext in ('png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'):
                context_lines.append(
                    f"IMAGE FILE at: {local_path}. Use analyze_image IMMEDIATELY with this path."
                )
            elif ext in ('mp3', 'wav', 'm4a', 'ogg', 'flac'):
                context_lines.append(
                    f"AUDIO FILE at: {local_path}. Use transcribe_audio IMMEDIATELY with this path."
                )
            else:
                context_lines.append(
                    f"DATA FILE at: {local_path}. Use read_file IMMEDIATELY with this path before anything else."
                )
        if youtube_match:
            context_lines.append(
                f"YouTube video: {youtube_match.group(0)}. Use get_youtube_transcript first."
            )
        context_note = "\n".join(context_lines)

        try:
            answer = run_react_agent(question, context_note=context_note)
        except Exception as e:
            print(f"Failed: {e}")
            answer = "unknown"

        print(f"Answer: {answer[:80]}")
        return answer

# =========================================================
# GRADIO APP (submission runner) — UI and eval-API contract unchanged.
# =========================================================
def run_and_submit_all(profile: gr.OAuthProfile | None):
    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"

    try:
        agent = BasicAgent(api_url=api_url)
    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(agent_code)

    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 "Fetched questions list is empty or invalid format.", None
        print(f"Fetched {len(questions_data)} questions.")
    except requests.exceptions.RequestException as e:
        print(f"Error fetching questions: {e}")
        return f"Error fetching questions: {e}", None
    except requests.exceptions.JSONDecodeError as e:
        print(f"Error decoding JSON response from questions endpoint: {e}")
        print(f"Response text: {response.text[:500]}")
        return f"Error decoding server response for questions: {e}", None
    except Exception as e:
        print(f"An unexpected error occurred fetching questions: {e}")
        return f"An unexpected error occurred fetching questions: {e}", None

    results_log = []
    answers_payload = []
    print(f"Running agent on {len(questions_data)} questions...")
    for item in questions_data:
        task_id = item.get("task_id")
        question_text = item.get("question")
        if not task_id or question_text is None:
            print(f"Skipping item with missing task_id or question: {item}")
            continue
        try:
            submitted_answer = agent(question_text, task_id=task_id)
            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})
        except Exception as e:
            print(f"Error running agent on task {task_id}: {e}")
            results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"})

    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)

    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)

    print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
    try:
        response = requests.post(submit_url, json=submission_data, timeout=120)
        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.")
        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}"
        print(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."
        print(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}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except Exception as e:
        status_message = f"An unexpected error occurred during submission: {e}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df


with gr.Blocks() as demo:
    gr.Markdown("# GAIA Agent Evaluation Runner (ReAct + Reflection)")
    gr.Markdown(
        """
        **Instructions:**
        1. Set `GOOGLE_API_KEY` and `GROQ_API_KEY` as secrets in your Space settings.
        2. Log in with the button below.
        3. Click 'Run Evaluation & Submit All Answers'.

        Every question runs through a single ReAct loop: the agent reasons, decides
        for itself whether it needs web_search / wikipedia_search / calculator /
        python_tool / read_file / analyze_image / transcribe_audio /
        get_youtube_transcript, observes the result, and repeats until confident —
        followed by a reflection pass that sends it back to research further if its
        own answer isn't fully supported by the evidence it gathered.
        """
    )
    gr.LoginButton()
    run_button = gr.Button("Run Evaluation & Submit All Answers")
    status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, 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])

if __name__ == "__main__":
    print("\n" + "-" * 30 + " App Starting " + "-" * 30)
    space_host_startup = os.getenv("SPACE_HOST")
    space_id_startup = os.getenv("SPACE_ID")

    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(f"✅ SPACE_ID found: {space_id_startup}")
        print(f"   Repo URL: https://huggingface.co/spaces/{space_id_startup}")
    else:
        print("ℹ️  SPACE_ID environment variable not found (running locally?).")

    print("-" * (60 + len(" App Starting ")) + "\n")
    print("Launching Gradio Interface for GAIA Agent Evaluation...")
    demo.launch(debug=True, share=False, ssr_mode=False)