import os
import tempfile
import time
from collections import deque
import gradio as gr
import requests
import inspect
import pandas as pd
import spaces
from huggingface_hub import hf_hub_download

from smolagents import (
    ActionStep,
    LiteLLMModel,
    PythonInterpreterTool,
    ToolCallingAgent,
    WebSearchTool,
    VisitWebpageTool,
    WikipediaSearchTool,
)

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

# The scoring server's own /files/{task_id} endpoint is broken (confirmed via
# a direct request and a matching open PR on agents-course/Unit4_scoring:
# it stores GAIA attachment paths as HF-Hub-repo-relative paths but checks
# them as local filesystem paths, so every lookup 404s: "No file path
# associated with task_id ..."). Attachments live directly in the gated
# gaia-benchmark/GAIA dataset instead; fetch them from there using the
# logged-in user's own OAuth token (requires that account to have requested
# access to the dataset, and the `gated-repos` OAuth scope in README.md).
GAIA_DATASET_REPO = "gaia-benchmark/GAIA"
GAIA_DATASET_SUBDIR = "2023/validation"


@spaces.GPU
def _zerogpu_startup_check():
    # This Space runs on ZeroGPU hardware but the agent below only makes
    # network calls (Groq API, web search) and never touches CUDA.
    # ZeroGPU requires at least one @spaces.GPU function to be declared,
    # so this no-op satisfies that check without spending any GPU quota
    # (it is never actually invoked).
    return None

# GAIA benchmark expects a terse, exact-match final answer.
GAIA_ANSWER_FORMAT_INSTRUCTIONS = """You are a general AI assistant. I will ask you a question.
Report your thoughts, and finish your work by calling final_answer() with your answer.
Your final 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 commas to write it, and don't use units such as $ or % unless specified otherwise.
If you are asked for a string, don't use articles or abbreviations (e.g. for cities), and write digits in plain text unless specified otherwise.
If you are asked for a comma separated list, apply the above rules to each element depending on whether it's a number or a string.
Work efficiently: if a search or lookup doesn't find what you need after 1-2 tries, try a meaningfully different approach rather than repeating similar queries, and give your best-guess final_answer rather than exhausting all steps.
Never call visit_webpage on a wikipedia.org URL - it always returns 403 Forbidden. Use the wikipedia_search tool for Wikipedia content instead.
"""

class TokenPacer:
    """
    Step_callback that tracks actual token usage per step in a trailing 60s
    window and sleeps as needed to stay under a tokens-per-minute budget.
    Necessary because a single call's fixed overhead (system prompt + tool
    schemas + question, before any tool output) already runs 2,400+ tokens
    on this agent and grows with context - smolagents' native
    requests_per_minute throttle can't account for that since it only paces
    call count, not size. This is a secondary safeguard: Cerebras' free-tier
    TPM budget (30,000) is generous enough that this should rarely trigger.
    """
    def __init__(self, tokens_per_minute_budget: int = 5000):
        self.tokens_per_minute_budget = tokens_per_minute_budget
        self._usage_window: deque[tuple[float, int]] = deque()

    def __call__(self, memory_step: ActionStep, agent: ToolCallingAgent) -> None:
        now = time.monotonic()
        usage = getattr(memory_step, "token_usage", None)
        tokens = (usage.input_tokens + usage.output_tokens) if usage else 0
        self._usage_window.append((now, tokens))
        cutoff = now - 60
        while self._usage_window and self._usage_window[0][0] < cutoff:
            self._usage_window.popleft()
        window_tokens = sum(t for _, t in self._usage_window)
        if window_tokens > self.tokens_per_minute_budget and self._usage_window:
            wait = 60 - (now - self._usage_window[0][0]) + 0.5
            if wait > 0:
                print(f"Token pacer: {window_tokens} tokens in the last 60s (budget {self.tokens_per_minute_budget}), sleeping {wait:.1f}s")
                time.sleep(wait)


class MemoryTrimmer:
    """
    Step_callback that collapses old tool outputs in the agent's memory.
    The agent re-sends the *entire* step history on every call, so input
    tokens grow every single step (2k -> 5k -> 8k -> ... -> 20k+ by step 6),
    which both wrecks the TPM budget and slows every later step far more
    than it needs to. Keeps the most recent `keep_recent` steps' tool
    outputs intact (the agent still needs that detail) and truncates older
    ones to a short placeholder, keeping per-call token cost roughly flat
    across a run instead of growing unbounded.
    """
    def __init__(self, keep_recent: int = 2, max_old_observation_chars: int = 300):
        self.keep_recent = keep_recent
        self.max_old_observation_chars = max_old_observation_chars

    def __call__(self, memory_step: ActionStep, agent: ToolCallingAgent) -> None:
        action_steps = [step for step in agent.memory.steps if isinstance(step, ActionStep)]
        for step in action_steps[:-self.keep_recent]:
            observations = getattr(step, "observations", None)
            if observations and len(observations) > self.max_old_observation_chars:
                step.observations = observations[:self.max_old_observation_chars] + " [...older output truncated to save context]"


# --- Basic Agent Definition ---
# ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
class BasicAgent:
    def __init__(self):
        # gpt-oss-120b's daily token quota (1M, separate per model on
        # Cerebras) got fully used up mid-session. zai-glm-4.7 has its own
        # untouched 1M/day quota, so switch there rather than waiting ~24h
        # for gpt-oss-120b's to reset.
        model_id = os.getenv("AGENT_MODEL_ID", "cerebras/zai-glm-4.7")
        api_key = os.getenv("CEREBRAS_API_KEY")
        if not api_key:
            print("Warning: CEREBRAS_API_KEY is not set - the agent will fail to call the model.")

        # Cerebras' free tier: 5 RPM, 30,000 TPM, 1,000,000 TPD - far more
        # TPM headroom than Groq's free tier (6000-12000), which was the
        # actual bottleneck there. requests_per_minute paces call frequency
        # natively; retry=False avoids smolagents' internal retry-on-error
        # firing extra invisible HTTP calls on a rate-limited response.
        self.model = LiteLLMModel(
            model_id=model_id,
            api_key=api_key,
            temperature=0,
            requests_per_minute=float(os.getenv("RATE_LIMIT_RPM", "4.5")),
            retry=False,
        )
        # ToolCallingAgent (not CodeAgent): gpt-oss-120b kept mixing reasoning
        # text into its code instead of cleanly wrapping it in <code> tags,
        # so CodeAgent's regex-based code-block parser failed on a large
        # fraction of steps - wasted steps that burned tokens/rate-limit
        # budget without making progress, and led to hallucinated final
        # answers. ToolCallingAgent uses the provider's native structured
        # tool-calling instead of parsing free-form text, which sidesteps
        # this failure mode entirely. PythonInterpreterTool replaces the
        # code-execution capability CodeAgent had built in.
        self.agent = ToolCallingAgent(
            model=self.model,
            tools=[
                WebSearchTool(),
                VisitWebpageTool(max_output_length=3000),  # default 40000 chars blows the TPM budget in one call
                WikipediaSearchTool(content_type="summary"),  # "text" (default) returns the full article
                PythonInterpreterTool(authorized_imports=[
                    "pandas", "numpy", "math", "re", "json", "itertools",
                    "collections", "statistics", "datetime", "io", "openpyxl", "PIL",
                ]),
            ],
            # Reverted 12 -> 7: raising it caused several questions to run
            # 12-13 steps each with growing context, and the cumulative
            # token usage blew Cerebras' *daily* quota partway through the
            # run (confirmed: "Tokens per day limit exceeded" starting
            # around question 13) - every question after that auto-failed,
            # dropping the score from 25% to 10%. 7 steps was empirically
            # better: fewer wasted tokens per question, more questions
            # actually get a real shot before the daily budget runs out.
            max_steps=7,
            step_callbacks=[
                MemoryTrimmer(),
                TokenPacer(tokens_per_minute_budget=int(os.getenv("RATE_LIMIT_TOKENS_PER_MINUTE", "25000"))),
            ],
        )
        print("BasicAgent initialized.")

    def __call__(self, question: str, file_path: str | None = None) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        task = GAIA_ANSWER_FORMAT_INSTRUCTIONS + f"\nQuestion: {question}"
        if file_path:
            task += (
                f"\n\nA file for this question was downloaded locally to: {file_path}\n"
                "Open/read it with Python (pandas, openpyxl, PIL, etc. as appropriate) to answer the question."
            )
        try:
            answer = self.agent.run(task)
        except Exception as e:
            print(f"Agent run failed: {e}")
            return f"AGENT ERROR: {e}"
        answer = str(answer).strip()
        print(f"Agent returning answer: {answer}")
        return answer

def run_and_submit_all( profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | 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"

    # 1. Instantiate Agent ( modify this part to create your agent)
    try:
        agent = BasicAgent()
    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)

    # 2. Fetch Questions
    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

    # 3. Run your Agent
    results_log = []
    answers_payload = []
    print(f"Running agent on {len(questions_data)} questions...")
    with tempfile.TemporaryDirectory() as tmp_dir:
        for item in questions_data:
            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:
                print(f"Skipping item with missing task_id or question: {item}")
                continue

            file_path = None
            if file_name:
                try:
                    file_path = hf_hub_download(
                        repo_id=GAIA_DATASET_REPO,
                        repo_type="dataset",
                        filename=f"{GAIA_DATASET_SUBDIR}/{file_name}",
                        token=oauth_token.token if oauth_token else None,
                    )
                except Exception as e:
                    print(f"Could not download attached file for task {task_id} from {GAIA_DATASET_REPO}: {e}")
                    # Fall back to the scoring server's own endpoint, in case
                    # it's since been fixed (see GAIA_DATASET_REPO comment above).
                    try:
                        file_response = requests.get(f"{api_url}/files/{task_id}", timeout=30)
                        file_response.raise_for_status()
                        file_path = os.path.join(tmp_dir, file_name)
                        with open(file_path, "wb") as f:
                            f.write(file_response.content)
                    except requests.exceptions.RequestException as e2:
                        print(f"Fallback download also failed for task {task_id}: {e2}")
                        file_path = None

            try:
                submitted_answer = agent(question_text, file_path=file_path)
                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)

    # 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=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.")
        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


# --- 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.

        **Setup:** This agent calls Cerebras (zai-glm-4.7) via `smolagents`. Get a free key at https://cloud.cerebras.ai and set it as the `CEREBRAS_API_KEY` secret in this Space's settings before running.
        """
    )

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

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 Basic Agent Evaluation...")
    demo.launch(debug=True, share=False)