import os
import re
import io
import time
import random
import base64
import threading
from typing import Optional, TypedDict

import requests
import pandas as pd
import gradio as gr
import spaces
from langgraph.graph import StateGraph, END


@spaces.GPU
def _zerogpu_warmup():
    # Dummy function so ZeroGPU hardware detects at least one @spaces.GPU
    # function at startup. Our agent doesn't actually need GPU compute
    # (Gemini runs remotely via API), this just satisfies the platform check.
    return True

# --- Constants ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
MAX_STEPS = 6  # max reasoning/tool-use loops per question (lower = fewer Gemini calls burned per question)

# Free-tier gemini-2.5-flash quota is tight (RPM-limited) and the *actual*
# observed ceiling is often lower than the documented one. Start conservative
# and let the adaptive limiter below back off further if we still get 429s.
MIN_SECONDS_BETWEEN_CALLS = 9.0
MAX_INTERVAL_BETWEEN_CALLS = 30.0
MAX_RETRIES = 7
COOLDOWN_AFTER_EXHAUSTED_RETRIES = 25.0  # extra pause before moving to the next question

SYSTEM_PROMPT = """You are a careful, precise general-purpose research assistant answering
GAIA benchmark questions. Your answer will be graded by EXACT STRING MATCH against a
ground truth, so formatting matters enormously.

You have access to these tools:
- search: web_search("your query") -> returns a few web search result snippets
- python: python_eval("expression or short script") -> returns stdout / result
- (A file may already be attached to this question below, already read for you.)

Work step by step. On each turn, reply with EXACTLY ONE of these two formats:

Thought: <your reasoning>
Action: search: <query>

OR

Thought: <your reasoning>
Action: python: <code>

When you are confident of the final answer, reply with ONLY:

Final Answer: <answer>

Formatting rules for the Final Answer (critical, exact match grading):
- If asked for a number: write digits only, no commas, no units unless explicitly asked
  for units, no $ sign unless asked.
- If asked for a string: no articles (a/an/the) unless part of a proper noun, no
  abbreviations unless asked, match the exact casing/spelling implied by the question.
- If asked for a comma separated list: apply the above rules to each element.
- Do not add explanations, punctuation, or extra words around the Final Answer.
- Never include the literal text "FINAL ANSWER" anywhere in the answer itself.
"""


# --- Rate limiting + retry helpers -------------------------------------

class AdaptiveRateLimiter:
    """Enforces spacing between calls, and *learns* the real limit at runtime:
    every 429 pushes the interval up (we were going too fast), and a streak of
    clean successes eases it back down toward the floor. This matters because
    the documented free-tier RPM is often optimistic vs. what you actually get."""

    def __init__(self, min_interval: float, max_interval: float):
        self.floor = min_interval
        self.ceiling = max_interval
        self.interval = min_interval
        self._lock = threading.Lock()
        self._last_call = 0.0
        self._success_streak = 0

    def wait(self):
        with self._lock:
            now = time.monotonic()
            elapsed = now - self._last_call
            if elapsed < self.interval:
                time.sleep(self.interval - elapsed)
            self._last_call = time.monotonic()

    def penalize(self):
        with self._lock:
            self._success_streak = 0
            self.interval = min(self.ceiling, self.interval * 1.6 + 1.0)
            print(f"[rate-limiter] backing off — new spacing: {self.interval:.1f}s")

    def reward(self):
        with self._lock:
            self._success_streak += 1
            if self._success_streak >= 4 and self.interval > self.floor:
                self.interval = max(self.floor, self.interval * 0.85)
                self._success_streak = 0
                print(f"[rate-limiter] easing up — new spacing: {self.interval:.1f}s")


_gemini_limiter = AdaptiveRateLimiter(MIN_SECONDS_BETWEEN_CALLS, MAX_INTERVAL_BETWEEN_CALLS)


class ModelNotFound(Exception):
    """404 on the model URL — wrong/unavailable model name for this key. Not retryable."""
    pass


class DailyQuotaExhausted(Exception):
    """Raised when the 429 is clearly a daily (RPD) quota exhaustion, not a
    transient per-minute throttle. Retrying within the same day cannot help."""
    pass


def _retry_after_seconds(resp) -> float:
    """Pull a Retry-After hint from headers or the Gemini error body, if present."""
    header_val = resp.headers.get("Retry-After")
    if header_val:
        try:
            return float(header_val)
        except ValueError:
            pass
    try:
        body = resp.json()
        for detail in body.get("error", {}).get("details", []):
            if "retryDelay" in detail:
                delay = detail["retryDelay"]
                return float(str(delay).rstrip("s"))
    except Exception:
        pass
    return None


def _is_daily_quota_exhausted(resp) -> bool:
    """Best-effort check of the 429 error body for a per-day quota violation
    (e.g. quotaId/quotaMetric containing 'PerDay'). If Google's error format
    doesn't match what we expect, this just returns False and normal retry
    logic takes over — it's a fast-path, not the only safety net."""
    try:
        body = resp.json()
        details = body.get("error", {}).get("details", [])
        for detail in details:
            for violation in detail.get("violations", []):
                blob = str(violation).lower()
                if "perday" in blob.replace(" ", ""):
                    return True
            blob = str(detail).lower()
            if "perday" in blob.replace(" ", ""):
                return True
    except Exception:
        pass
    return False


def call_with_retry(fn, *, limiter: AdaptiveRateLimiter = None, max_retries: int = MAX_RETRIES):
    """Runs fn() with adaptive rate limiting + exponential backoff on 429 / 5xx.
    On 429s specifically, tells the limiter to slow down going forward so later
    calls (including the *next question's*) don't immediately hit the same wall."""
    last_exc = None
    for attempt in range(max_retries):
        if limiter:
            limiter.wait()
        try:
            resp = fn()
            if resp.status_code == 404:
                raise ModelNotFound(
                    f"404 Not Found for {getattr(resp, 'url', 'model URL')} — the model name is "
                    f"wrong or unavailable for this key. Retrying won't fix this."
                )
            if resp.status_code == 429 or resp.status_code >= 500:
                if resp.status_code == 429 and _is_daily_quota_exhausted(resp):
                    raise DailyQuotaExhausted(
                        "Daily request quota exhausted for this Gemini API key/project. "
                        "Retrying will not help until it resets (midnight Pacific time)."
                    )
                if resp.status_code == 429 and limiter:
                    limiter.penalize()
                wait_s = _retry_after_seconds(resp)
                if wait_s is None:
                    wait_s = (2 ** attempt) + random.uniform(0, 1.5)
                wait_s = min(wait_s, 60)
                print(f"[retry] status={resp.status_code}, backing off {wait_s:.1f}s "
                      f"(attempt {attempt + 1}/{max_retries})")
                time.sleep(wait_s)
                last_exc = requests.HTTPError(f"{resp.status_code} after retries")
                continue
            resp.raise_for_status()
            if limiter:
                limiter.reward()
            return resp
        except requests.RequestException as e:
            last_exc = e
            wait_s = (2 ** attempt) + random.uniform(0, 1.5)
            print(f"[retry] request error: {e}; backing off {wait_s:.1f}s "
                  f"(attempt {attempt + 1}/{max_retries})")
            time.sleep(wait_s)

    # Every retry failed — this usually means we're sustained-throttled, not just
    # unlucky. Pause further before returning control, so the *next* question
    # doesn't immediately walk into the same wall this one just hit.
    if limiter:
        print(f"[retry] exhausted, cooling down {COOLDOWN_AFTER_EXHAUSTED_RETRIES:.0f}s before continuing")
        time.sleep(COOLDOWN_AFTER_EXHAUSTED_RETRIES)
    raise last_exc or RuntimeError("call_with_retry: exhausted retries")


def web_search(query: str) -> str:
    def _do():
        from ddgs import DDGS
        # ddgs raises its own exceptions rather than returning a Response,
        # so wrap it in a tiny shim compatible with call_with_retry.
        class _FakeResp:
            status_code = 200
            def raise_for_status(self):
                pass
        try:
            with DDGS() as ddgs:
                results = list(ddgs.text(query, max_results=5))
            _FakeResp._results = results
            return _FakeResp()
        except Exception as e:
            msg = str(e).lower()
            fr = _FakeResp()
            fr.status_code = 429 if "ratelimit" in msg or "429" in msg else 500
            fr._error = str(e)
            return fr

    try:
        resp = call_with_retry(_do, limiter=None, max_retries=3)
        results = getattr(resp, "_results", None)
        if not results:
            return "No results found."
        out = []
        for r in results:
            out.append(f"- {r.get('title','')}: {r.get('body','')} ({r.get('href','')})")
        return "\n".join(out)
    except Exception as e:
        return f"Search error: {e}"


def python_eval(code: str) -> str:
    import contextlib
    buf = io.StringIO()
    local_vars = {}
    try:
        with contextlib.redirect_stdout(buf):
            try:
                result = eval(code, {}, local_vars)
                if result is not None:
                    print(result)
            except SyntaxError:
                exec(code, {}, local_vars)
        output = buf.getvalue().strip()
        return output if output else "(no output)"
    except Exception as e:
        return f"Python error: {e}"


AUDIO_EXTENSIONS = {".mp3": "audio/mpeg", ".wav": "audio/wav", ".m4a": "audio/mp4"}
SPREADSHEET_EXTENSIONS = (".xlsx", ".xls")

YOUTUBE_RE = re.compile(
    r"https?://(?:www\.)?(?:youtube\.com/watch\?v=[\w-]+|youtu\.be/[\w-]+)"
)


def _spreadsheet_to_text(raw_bytes: bytes) -> str:
    """Parses an .xlsx/.xls file into a plain-text table the model can read directly,
    instead of leaving it as opaque binary the agent has no way to inspect."""
    try:
        sheets = pd.read_excel(io.BytesIO(raw_bytes), sheet_name=None)
    except Exception as e:
        return f"[Could not parse spreadsheet: {e}]"
    chunks = []
    for name, df in sheets.items():
        chunks.append(f"--- Sheet: {name} ---\n{df.to_csv(index=False)}")
    text = "\n\n".join(chunks)
    if len(text) > 8000:
        text = text[:8000] + "\n...[truncated]"
    return text


def download_task_file(api_url: str, task_id: str):
    """Returns (kind, content, mime) where kind is 'image', 'audio', 'text', 'binary', or None."""
    try:
        resp = requests.get(f"{api_url}/files/{task_id}", timeout=30)
        if resp.status_code != 200:
            return None, None, None
        content_type = resp.headers.get("content-type", "")
        disposition = resp.headers.get("content-disposition", "")
        fname_match = re.search(r'filename="?([^";]+)"?', disposition)
        fname = fname_match.group(1).lower() if fname_match else ""

        if "image" in content_type:
            return "image", resp.content, content_type

        if "audio" in content_type:
            return "audio", resp.content, content_type

        # Fall back to filename extension for audio if content-type is generic
        # (e.g. application/octet-stream, which HF file endpoints sometimes use)
        for ext, mime in AUDIO_EXTENSIONS.items():
            if fname.endswith(ext):
                return "audio", resp.content, mime

        if "spreadsheet" in content_type or fname.endswith(SPREADSHEET_EXTENSIONS):
            return "text", _spreadsheet_to_text(resp.content), "text/csv"

        try:
            text = resp.content.decode("utf-8")
            if len(text) > 6000:
                text = text[:6000] + "\n...[truncated]"
            return "text", text, "text/plain"
        except UnicodeDecodeError:
            return "binary", None, content_type
    except Exception:
        return None, None, None


# If GEMINI_MODEL is set as a Space secret/variable, that's tried first.
# Otherwise we probe this list in order and use the first one that doesn't 404
# for this key/project. Free-tier default models have shifted more than once
# in 2026 (2.5 Flash -> 3 Flash), so hardcoding one name is fragile.
MODEL_CANDIDATES = [m for m in [
    os.getenv("GEMINI_MODEL"),
    "gemini-2.5-flash",
    "gemini-flash-latest",
    "gemini-3-flash",
    "gemini-2.0-flash",
] if m]

GEMINI_API_BASE = "https://generativelanguage.googleapis.com/v1beta/models"


# --- LangGraph agent state + graph -------------------------------------
#
# This replaces the previous hand-rolled `for step in range(MAX_STEPS): ...`
# loop with an explicit LangGraph StateGraph. Same ReAct-style
# "Thought / Action / Observation / Final Answer" protocol and the same
# underlying GeminiAgent._call_model (with its adaptive rate limiter, retry
# logic, and daily-quota detection) — just orchestrated as graph nodes +
# conditional edges instead of a raw loop, so it's a proper LangGraph agent.

class AgentState(TypedDict):
    transcript: str
    step: int
    media_bytes: Optional[bytes]
    media_mime: Optional[str]
    youtube_url: Optional[str]
    pending_tool: Optional[str]
    pending_arg: Optional[str]
    final_answer: Optional[str]


def build_agent_graph(client: "GeminiAgent"):
    """Builds the LangGraph StateGraph for the ReAct loop:

        call_model --(final answer)--> END
        call_model --(tool call)-----> run_tool --> call_model
        call_model --(malformed, retry left)--> call_model
        call_model --(malformed, no retries left / step limit)--> END
    """

    def call_model(state: AgentState) -> AgentState:
        step = state["step"]
        if step >= MAX_STEPS:
            return {**state, "final_answer": "Unable to determine answer within step limit."}

        reply = client._call_model(
            state["transcript"],
            media_bytes=state["media_bytes"] if step == 0 else None,
            media_mime=state["media_mime"] if step == 0 else None,
            youtube_url=state["youtube_url"] if step == 0 else None,
        ).strip()
        print(f"[step {step}] model reply: {reply[:200]}")

        if reply.startswith("[MODEL_CALL_FAILED]"):
            return {**state, "final_answer": f"AGENT ERROR: {reply}", "step": step + 1}

        final_match = re.search(r"Final Answer:\s*(.+)", reply, re.DOTALL)
        if final_match:
            answer = final_match.group(1).strip()
            answer = re.sub(r"^FINAL ANSWER[:\s]*", "", answer, flags=re.IGNORECASE)
            return {**state, "final_answer": answer, "step": step + 1}

        action_match = re.search(r"Action:\s*(search|python):\s*(.+)", reply, re.DOTALL)
        if action_match:
            tool, arg = action_match.group(1), action_match.group(2).strip()
            return {
                **state,
                "transcript": state["transcript"] + f"\n{reply}",
                "pending_tool": tool,
                "pending_arg": arg,
                "step": step + 1,
            }

        # Reply matched neither format (e.g. got cut off, or the model got
        # confused by video/audio input). Nudge it back on format instead of
        # silently submitting a fragment like "Thought:" as the final answer.
        if reply and step < MAX_STEPS - 1:
            transcript = state["transcript"] + (
                f"\n{reply}\nObservation: Your reply didn't match the required "
                f"'Action: search|python: ...' or 'Final Answer: ...' format. "
                f"Reply again using exactly one of those two formats.\n"
            )
            return {**state, "transcript": transcript, "step": step + 1}

        fallback = reply.split("\n")[0].strip() if reply else "Unable to determine answer."
        return {**state, "final_answer": fallback, "step": step + 1}

    def run_tool(state: AgentState) -> AgentState:
        tool, arg = state["pending_tool"], state["pending_arg"]
        observation = web_search(arg) if tool == "search" else python_eval(arg)
        transcript = state["transcript"] + f"\nObservation: {observation}\n"
        return {**state, "transcript": transcript, "pending_tool": None, "pending_arg": None}

    def route_after_model(state: AgentState) -> str:
        if state.get("final_answer") is not None:
            return END
        if state.get("pending_tool"):
            return "run_tool"
        return "call_model"

    graph = StateGraph(AgentState)
    graph.add_node("call_model", call_model)
    graph.add_node("run_tool", run_tool)
    graph.set_entry_point("call_model")
    graph.add_conditional_edges(
        "call_model",
        route_after_model,
        {"run_tool": "run_tool", "call_model": "call_model", END: END},
    )
    graph.add_edge("run_tool", "call_model")
    return graph.compile()


class GeminiAgent:
    def __init__(self):
        self.api_key = os.getenv("GEMINI_API_KEY")
        if not self.api_key:
            raise ValueError("GEMINI_API_KEY secret not set in this Space.")
        self.model = self._discover_model()
        self.graph = build_agent_graph(self)
        print(f"GeminiAgent initialized with model: {self.model}")

    def _discover_model(self) -> str:
        """Tries each candidate model name with a trivial request and picks the
        first one that isn't a 404 for this key. A 429 still counts as 'found' —
        it means the model exists but is rate-limited, which is a separate problem."""
        last_error = None
        for name in MODEL_CANDIDATES:
            url = f"{GEMINI_API_BASE}/{name}:generateContent"
            try:
                resp = requests.post(
                    url,
                    params={"key": self.api_key},
                    json={
                        "contents": [{"role": "user", "parts": [{"text": "hi"}]}],
                        "generationConfig": {"temperature": 0, "maxOutputTokens": 1},
                    },
                    timeout=20,
                )
            except requests.RequestException as e:
                last_error = e
                continue
            if resp.status_code == 404:
                print(f"[model-discovery] {name} -> 404, trying next candidate")
                last_error = f"{name}: 404 Not Found"
                continue
            # Any other response (200, 429, etc.) means this model name is valid
            # for this key — 429 is a rate-limit problem, not a wrong-model problem.
            print(f"[model-discovery] using model: {name} (probe status {resp.status_code})")
            return name
        raise ValueError(
            f"None of the candidate Gemini models are reachable with this API key "
            f"(tried: {MODEL_CANDIDATES}). Last error: {last_error}. "
            f"Check https://aistudio.google.com/apikey for which models your key/project can access."
        )

    def _call_model(self, text_prompt: str, media_bytes: bytes = None,
                     media_mime: str = None, youtube_url: str = None) -> str:
        parts = [{"text": text_prompt}]

        if media_bytes and media_mime:
            parts.append({
                "inline_data": {
                    "mime_type": media_mime,
                    "data": base64.b64encode(media_bytes).decode("utf-8"),
                }
            })

        if youtube_url:
            # Gemini 2.5 models can ingest public YouTube URLs directly for
            # native video/audio understanding (no download needed).
            parts.append({
                "file_data": {"file_uri": youtube_url}
            })

        payload = {
            "system_instruction": {"parts": [{"text": SYSTEM_PROMPT}]},
            "contents": [{"role": "user", "parts": parts}],
            "generationConfig": {"temperature": 0},
        }

        def _do():
            return requests.post(
                f"{GEMINI_API_BASE}/{self.model}:generateContent",
                params={"key": self.api_key},
                json=payload,
                timeout=90,
            )

        try:
            resp = call_with_retry(_do, limiter=_gemini_limiter)
        except (DailyQuotaExhausted, ModelNotFound):
            raise
        except Exception as e:
            return f"[MODEL_CALL_FAILED] {e}"

        data = resp.json()
        try:
            candidate_parts = data["candidates"][0]["content"]["parts"]
            return "".join(p.get("text", "") for p in candidate_parts)
        except (KeyError, IndexError):
            return ""

    def __call__(self, question: str, api_url: str = DEFAULT_API_URL, task_id: str = None) -> str:
        media_bytes, media_mime, youtube_url = None, None, None
        transcript = f"Question: {question}"

        yt_match = YOUTUBE_RE.search(question)
        if yt_match:
            youtube_url = yt_match.group(0)

        if task_id:
            kind, content, mime = download_task_file(api_url, task_id)
            if kind == "text":
                transcript += f"\n\n[Attached file content]:\n{content}"
            elif kind == "image":
                media_bytes, media_mime = content, mime
            elif kind == "audio":
                media_bytes, media_mime = content, mime
                transcript += "\n\n[An audio recording is attached — listen to it to answer.]"

        initial_state: AgentState = {
            "transcript": transcript,
            "step": 0,
            "media_bytes": media_bytes,
            "media_mime": media_mime,
            "youtube_url": youtube_url,
            "pending_tool": None,
            "pending_arg": None,
            "final_answer": None,
        }

        # recursion_limit is generous relative to MAX_STEPS since each
        # model step + tool call is 2 graph-node transitions.
        final_state = self.graph.invoke(
            initial_state, config={"recursion_limit": MAX_STEPS * 3 + 10}
        )

        if final_state.get("final_answer") and final_state["final_answer"].startswith("AGENT ERROR:"):
            return final_state["final_answer"]

        return final_state.get("final_answer") or "Unable to determine answer within step limit."


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:
        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 = GeminiAgent()
    except Exception as e:
        return f"Error initializing agent: {e}", None

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

    try:
        response = requests.get(questions_url, timeout=15)
        response.raise_for_status()
        questions_data = response.json()
        if not questions_data:
            return "Fetched questions list is empty or invalid format.", None
    except Exception as e:
        return f"Error fetching questions: {e}", None

    results_log = []
    answers_payload = []
    quota_exhausted = False
    consecutive_failures = 0
    CONSECUTIVE_FAILURE_LIMIT = 3  # fallback circuit breaker if quota detection doesn't fire

    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:
            continue
        try:
            submitted_answer = agent(question_text, api_url=api_url, task_id=task_id)
        except (DailyQuotaExhausted, ModelNotFound) as e:
            print(f"[circuit-breaker] stopping run: {e}")
            quota_exhausted = True
            break
        except Exception as e:
            submitted_answer = f"AGENT ERROR: {e}"

        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})

        if isinstance(submitted_answer, str) and submitted_answer.startswith("AGENT ERROR"):
            consecutive_failures += 1
            if consecutive_failures >= CONSECUTIVE_FAILURE_LIMIT:
                print(f"[circuit-breaker] {consecutive_failures} consecutive failures, stopping run early")
                quota_exhausted = True
                break
        else:
            consecutive_failures = 0

    if not answers_payload:
        msg = "Agent did not produce any answers to submit."
        if quota_exhausted:
            msg += (" Stopped early — this looks like your Gemini API quota is exhausted "
                     "(check https://aistudio.google.com/rate-limit). RPD quotas reset at "
                     "midnight Pacific time.")
        return msg, pd.DataFrame(results_log)

    submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
    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.')}"
        )
        if quota_exhausted:
            final_status += (
                f"\n\nNote: run stopped early after {len(answers_payload)} questions — Gemini API "
                f"quota appears exhausted. Check https://aistudio.google.com/rate-limit and try again "
                f"after it resets (RPD resets at midnight Pacific time)."
            )
        return final_status, pd.DataFrame(results_log)
    except Exception as e:
        return f"Submission Failed: {e}", pd.DataFrame(results_log)


with gr.Blocks() as demo:
    gr.Markdown("# GAIA Agent Evaluation Runner (Gemini-powered)")
    gr.Markdown(
        """
        1. Log in with your Hugging Face account below.
        2. Click 'Run Evaluation & Submit All Answers'.
        3. Wait — the agent paces its requests and backs off automatically if it
           gets rate-limited, so a full run can take 15-30+ minutes. Don't refresh mid-run.
        """
    )
    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)
    demo.launch(debug=True, share=False)