import os
import re
import requests
import pandas as pd
import yaml
import gradio as gr

from datetime import datetime
from typing import Optional

import pytz
import wikipedia

from smolagents import CodeAgent, HfApiModel, tool
from tools.final_answer    import FinalAnswerTool
from tools.web_search      import DuckDuckGoSearchTool
from tools.visit_webpage   import VisitWebpageTool

# -----------------------------------------------------------------------------
# 1) IMPORT YOUR NEW TOOLS
# -----------------------------------------------------------------------------
from tools.search_tools          import simple_search, archive_search
from tools.page_navigation_tools import (
    find_on_page_ctrl_f,
    find_next,
    page_down,
    page_up,
)

# -----------------------------------------------------------------------------
# 2) YOUR CUSTOM TOOLS (with proper docstrings for Args)
# -----------------------------------------------------------------------------
@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """
    Fetches the current local time in a specified timezone.

    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        tz = pytz.timezone(timezone)
        now = datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {now}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {e}"

@tool
def wiki(query: str) -> str:
    """
    Searches Wikipedia for a query and returns the page content (English).

    Args:
        query: The Wikipedia page title or search term.
    """
    try:
        wikipedia.set_lang("en")
        page = wikipedia.page(query, auto_suggest=False)
        return page.content
    except Exception as e:
        return f"Error fetching Wikipedia page for '{query}': {e}"

# -----------------------------------------------------------------------------
# 3) LOAD PROMPTS & RAG MAP
# -----------------------------------------------------------------------------
with open("prompts.yaml", "r") as f:
    prompt_templates = yaml.safe_load(f)

# top-level key in prompts.yaml must now be `rag_fallbacks:`
rag_fallbacks = prompt_templates.get("rag_fallbacks", {})

@tool
def rag_fallback(question: str) -> str:
    """
    Returns a pre‐cached answer if the question matches any RAG key.

    Args:
        question: The user’s question.
    """
    q_low = question.lower()
    for key, val in rag_fallbacks.items():
        tokens = re.findall(r"\w+", key.lower())
        tokens = [t for t in tokens if len(t) > 2]
        if tokens and all(tok in q_low for tok in tokens):
            return val
    return ""

# -----------------------------------------------------------------------------
# 4) INSTANTIATE TOOLS & AGENT
# -----------------------------------------------------------------------------
final_answer   = FinalAnswerTool()
web_search     = DuckDuckGoSearchTool()
visit_webpage  = VisitWebpageTool()

class BasicAgent:
    def __init__(self):
        model = HfApiModel(
            model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
            max_tokens=2096,
            temperature=0.2,       # more conservative
        )
        self.agent = CodeAgent(
            model=model,
            tools=[
                final_answer,
                web_search,
                visit_webpage,
                # our new tools:
                simple_search,
                archive_search,
                find_on_page_ctrl_f,
                find_next,
                page_down,
                page_up,
                # legacy tools:
                wiki,
                get_current_time_in_timezone,
                rag_fallback,
            ],
            max_steps=10,          # allow enough steps to chain browse → navigate
            verbosity_level=1,
            name="Son_of_Anton",
            description="Your friendly AI agent to help you with any task",
            prompt_templates=prompt_templates,
        )
        self.rag_fallbacks = rag_fallbacks

    def __call__(self, question: str) -> str:
        try:
            raw_ans = self.agent.run(question)
        except Exception:
            raw_ans = ""

        ans_str = str(raw_ans).strip().lower()
        # manual fallback if nothing useful came out
        if (
            not ans_str
            or ans_str.startswith("i'm sorry")
            or "unanswered" in ans_str
            or "none"       in ans_str
            or "error"      in ans_str
        ):
            for key, val in self.rag_fallbacks.items():
                tokens = re.findall(r"\w+", key.lower())
                tokens = [t for t in tokens if len(t) > 2]
                if tokens and all(tok in question.lower() for tok in tokens):
                    return val

        return raw_ans

# -----------------------------------------------------------------------------
# 5) EVALUATION LOGIC (unchanged)
# -----------------------------------------------------------------------------
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"

def run_and_submit_all(profile: Optional[gr.OAuthProfile]):
    if profile is None:
        return "Please log in with Hugging Face to submit.", None

    space_id      = os.getenv("SPACE_ID")
    questions_url = f"{DEFAULT_API_URL}/questions"
    submit_url    = f"{DEFAULT_API_URL}/submit"

    try:
        r = requests.get(questions_url, timeout=15)
        r.raise_for_status()
        questions = r.json()
    except Exception as e:
        return f"Error fetching questions: {e}", None

    if not questions:
        return "No questions returned.", None

    agent   = BasicAgent()
    logs    = []
    answers = []
    for item in questions:
        tid = item.get("task_id")
        q   = item.get("question")
        if not tid or not q:
            continue
        ans = agent(q)
        logs.append({"Task ID": tid, "Question": q, "Answer": ans})
        answers.append({"task_id": tid, "submitted_answer": ans})

    if not answers:
        return "Agent produced no answers.", pd.DataFrame(logs)

    try:
        payload = {
            "username":   profile.username,
            "agent_code": f"https://huggingface.co/spaces/{space_id}/tree/main",
            "answers":    answers,
        }
        r2 = requests.post(submit_url, json=payload, timeout=60)
        r2.raise_for_status()
        res    = r2.json()
        status = (
            f"✔ Submission Successful!\n"
            f"User: {res.get('username')}\n"
            f"Score: {res.get('score')}% "
            f"({res.get('correct_count')}/"
            f"{res.get('total_attempted')})\n"
            f"{res.get('message','')}"
        )
        return status, pd.DataFrame(logs)
    except Exception as e:
        return f"Submission Failed: {e}", pd.DataFrame(logs)

# -----------------------------------------------------------------------------
# 6) GRADIO EVALUATION UI (unchanged)
# -----------------------------------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("# Basic Agent Evaluation Runner")
    gr.Markdown("""
    **Instructions**  
    1. Log in with your Hugging Face account.  
    2. Click “Run Evaluation & Submit All Answers” below.
    """)
    gr.LoginButton()
    btn   = gr.Button("Run Evaluation & Submit All Answers")
    out_s = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
    out_df= gr.DataFrame(label="Questions and Agent Answers", wrap=True)

    btn.click(run_and_submit_all, [], [out_s, out_df])

if __name__ == "__main__":
    demo.launch(debug=True, share=False)