import os
import gradio as gr
import requests
import inspect
import pandas as pd
from smolagents import OpenAIServerModel, DuckDuckGoSearchTool, CodeAgent, WikipediaSearchTool
from smolagents.tools import Tool
import time
import openai
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
import random
import re
from collections import Counter

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

# --- Basic Agent Definition ---
# ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------


class TextSummarizationTool(Tool):
    """Summarizes a long text into a concise version by extracting leading sentences."""

    name = "text_summarization"
    description = "Summarizes a long input text into a short paragraph."

    inputs = {
        "text": {
            "type": "string",
            "description": "The long text that needs to be summarized.",
        }
    }
    output_type = "string"

    def forward(self, text: str) -> str:
        try:
            sentences = text.split('. ')
            if len(sentences) <= 3:
                return text
            return '. '.join(sentences[:3]) + '.'
        except Exception as e:
            return f"Error summarizing text: {e}"

class KeywordExtractorTool(Tool):
    """Extracts top keywords from a given block of text based on frequency."""

    name = "keyword_extractor"
    description = "Extracts the most frequent keywords from the provided text."

    inputs = {
        "text": {
            "type": "string",
            "description": "The text to analyze for keywords.",
        }
    }
    output_type = "string"

    def forward(self, text: str) -> str:
        try:
            words = re.findall(r'\b\w+\b', text.lower())
            stop_words = {'the', 'and', 'is', 'in', 'it', 'of', 'to', 'a'}
            filtered_words = [w for w in words if w not in stop_words]
            word_counts = Counter(filtered_words)
            keywords = ', '.join(word for word, _ in word_counts.most_common(5))
            return keywords
        except Exception as e:
            return f"Error extracting keywords: {e}"

class TextTranslationTool(Tool):
    """Translates simple words from source to target language using a dictionary lookup."""

    name = "text_translation"
    description = "Translates simple words from English to Spanish using a fixed dictionary."

    inputs = {
        "text": {
            "type": "string",
            "description": "The text to translate word-by-word.",
        },
        "source_lang": {
            "type": "string",
            "description": "Source language code (e.g., 'en').",
        },
        "target_lang": {
            "type": "string",
            "description": "Target language code (e.g., 'es').",
        }
    }
    output_type = "string"

    def __init__(self):
        self.translation_dict = {
            'hello': 'hola',
            'world': 'mundo',
            'goodbye': 'adiós',
            'thank': 'gracias',
            'you': 'tú'
        }

    def forward(self, text: str, source_lang: str, target_lang: str) -> str:
        try:
            words = text.split()
            translated_words = [self.translation_dict.get(word.lower(), word) for word in words]
            return ' '.join(translated_words)
        except Exception as e:
            return f"Error translating text: {e}"

# --- Retry Helper for Agent Call ---
def safe_agent_call(agent, question, retries=5, wait_time=20):
    """
    Helper function to safely call the agent with retry on rate limit errors (HTTP 429).
    """
    for attempt in range(retries):
        try:
            return agent(question)
        except Exception as e:
            error_text = str(e).lower()
            if "rate limit" in error_text or "429" in error_text:
                print(f"[Retry] Rate limit hit. Waiting {wait_time} seconds before retrying... (Attempt {attempt + 1}/{retries})")
                time.sleep(wait_time)
            else:
                print(f"[Error] Non-rate-limit error encountered: {e}")
                raise e
    raise Exception(f"Failed after {retries} retries due to repeated rate limit errors.")


# --- Basic Agent Definition ---
class BasicAgent:
    def __init__(self):
        self.agent = CodeAgent(
            model=OpenAIServerModel(model_id="gpt-4.1-mini"),
            tools=[
                DuckDuckGoSearchTool(),
                WikipediaSearchTool(),
                KeywordExtractorTool(),
                TextSummarizationTool(),
                TextTranslationTool()
            ],
            add_base_tools=True,
        )
        print("✅ BasicAgent initialized.")

    def __call__(self, question: str) -> str:
        """
        Calls the agent's run method to generate a response to the question.
        """
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        fixed_answer = self.agent.run(question)
        print(f"Agent returning answer: {fixed_answer}")
        return fixed_answer

# --- Main Logic for Fetching Questions, Running Agent, Submitting Answers ---
def run_and_submit_all(profile: gr.OAuthProfile | None):
    """
    Fetches all questions, runs BasicAgent on them with retry logic on rate limit,
    submits all answers, and displays the results.
    """

    # Determine HF Space runtime info
    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"

    # Instantiate Agent
    try:
        agent = BasicAgent()
    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(f"Agent Code Repository: {agent_code}")

    # 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 Exception as e:
        print(f"Error fetching questions: {e}")
        return f"Error fetching questions: {e}", None

    # Run Agent on Questions
    results_log = []
    answers_payload = []
    print(f"Running agent on {len(questions_data)} questions...")
    for idx, item in enumerate(questions_data, start=1):
        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)
            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}"})
    
        print(f"[{idx}/{len(questions_data)}] Waiting 60 seconds before next request...")
        time.sleep(60)

    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)

    # Prepare Submission
    submission_data = {
        "username": username.strip(),
        "agent_code": agent_code,
        "answers": answers_payload,
    }
    print(f"Submitting {len(answers_payload)} answers...")

    # Submit Answers
    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.')}"
        )
        results_df = pd.DataFrame(results_log)
        return final_status, results_df
    except Exception as e:
        print(f"Submission error: {e}")
        results_df = pd.DataFrame(results_log)
        return f"Submission Failed: {e}", 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.
        """
    )

    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)