import os
import re
import requests
import time
import gradio as ui
from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    VisitWebpageTool,
    InferenceClientModel,  # Reverted to the universally supported base class
    tool,
)

# ─────────────────────────────────────────────
# CONFIG
# ─────────────────────────────────────────────
API_BASE_URL = "https://agents-course-unit4-scoring.hf.space"

# ─────────────────────────────────────────────
# CUSTOM TOOLS  (pure-requests, zero extra deps)
# ─────────────────────────────────────────────

@tool
def wikipedia_search(query: str) -> str:
    """Searches Wikipedia for a topic and returns the article introduction.
    Use this to look up facts about people, places, events, and concepts.

    Args:
        query: The topic or person to search for on Wikipedia.

    Returns:
        The introductory text of the best matching Wikipedia article.
    """
    headers = {"User-Agent": "GAIA-agent/1.0 (huggingface-course)"}
    api = "https://en.wikipedia.org/w/api.php"
    try:
        r = requests.get(api, params={
            "action": "query", "list": "search",
            "srsearch": query, "format": "json", "srlimit": 1,
        }, headers=headers, timeout=15)
        r.raise_for_status()
        results = r.json().get("query", {}).get("search", [])
        if not results:
            return f"No Wikipedia article found for: {query}"
        title = results[0]["title"]

        r2 = requests.get(api, params={
            "action": "query", "titles": title,
            "prop": "extracts", "exintro": True,
            "explaintext": True, "format": "json",
        }, headers=headers, timeout=15)
        r2.raise_for_status()
        pages = r2.json().get("query", {}).get("pages", {})
        extract = next(iter(pages.values())).get("extract", "").strip()
        return f"Wikipedia — {title}:\n{extract[:4000]}" if extract else f"No extract for '{title}'"
    except Exception as e:
        return f"Wikipedia error: {e}"


@tool
def read_local_file(file_path: str) -> str:
    """Reads a local file and returns its content as text.
    Handles .txt, .csv, .json, .md, .html, .xml, .xlsx, .pdf

    Args:
        file_path: Path to the file on disk.

    Returns:
        The file contents as a string.
    """
    if not os.path.exists(file_path):
        return f"ERROR: File not found: {file_path}"
    ext = os.path.splitext(file_path)[1].lower()
    try:
        if ext in (".txt", ".md", ".py", ".html", ".xml", ".csv", ".json"):
            with open(file_path, encoding="utf-8", errors="replace") as f:
                return f.read()[:8000]
        elif ext == ".xlsx":
            import openpyxl
            wb = openpyxl.load_workbook(file_path, data_only=True)
            lines = []
            for name in wb.sheetnames:
                lines.append(f"=== {name} ===")
                for row in wb[name].iter_rows(values_only=True):
                    lines.append("\t".join("" if c is None else str(c) for c in row))
            return "\n".join(lines)[:8000]
        elif ext == ".pdf":
            from pypdf import PdfReader
            pages = PdfReader(file_path).pages
            return "\n".join(p.extract_text() or "" for p in pages)[:8000]
        else:
            with open(file_path, encoding="utf-8", errors="replace") as f:
                return f.read()[:8000]
    except Exception as e:
        return f"ERROR reading file: {e}"


# ─────────────────────────────────────────────
# AGENT SETUP
# ─────────────────────────────────────────────

INSTRUCTIONS = """You are an elite, highly precise research assistant solving GAIA benchmark tasks.

STRICT ANSWER FORMAT RULES:
- Your final output must contain ONLY the raw target answer value string — absolutely nothing else.
- Never add introductory phrases ("The answer is", "Result:"), markdown blocks, or finishing periods.
- For lists, use simple comma separation: item1, item2, item3

TACTICAL STRATEGY FOR WEB RESEARCH:
1. If you try to visit a webpage using 'visit_webpage' and receive an HTTP 403 Forbidden error, DO NOT keep trying to visit that site or similar sites.
2. Instead, rely on 'wikipedia_search' for any entity lookups or historical information.
3. If you need details from a specific platform that blocks you, run a targeted query using 'web_search', such as: web_search(query="site:en.wikipedia.org Mercedes Sosa studio albums 2000..2009 list"). Text snippets returned inside search results logs contain the exact numbers you need! Read those snippets carefully.
"""

# Restored InferenceClientModel with token support now that limits are fully reset
model = InferenceClientModel(
    model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
    token=os.environ.get("HF_TOKEN")
)

agent = CodeAgent(
    model=model,
    tools=[
        DuckDuckGoSearchTool(),
        VisitWebpageTool(),
        wikipedia_search,
        read_local_file,
    ],
    max_steps=6,
    verbosity_level=1,
    instructions=INSTRUCTIONS,
    additional_authorized_imports=["pandas", "numpy", "requests", "bs4", "openpyxl", "pypdf", "math", "re"]
)

# ─────────────────────────────────────────────
# ANSWER CLEANING
# ─────────────────────────────────────────────

def clean_answer(raw: str) -> str:
    ans = str(raw).strip()
    ans = re.sub(
        r"(?i)^(final\s+answer|the\s+(final\s+)?answer\s+(is\s*)?|answer\s*[:：]?\s*|result\s*[:：]?\s*|output\s*[:：]?\s*)",
        "", ans,
    ).strip().lstrip(":：").strip()

    first_block = ans.split("\n\n")[0].strip()
    if len(first_block) < 200:
        ans = first_block

    lines = [l.strip() for l in ans.splitlines() if l.strip()]
    if lines and len(lines[0]) < 200:
        ans = lines[0]

    ans = re.sub(r"^[\*\_`]+|[\*\_`]+$", "", ans).strip()
    if len(ans) >= 2 and ans[0] in ('"', "'") and ans[0] == ans[-1]:
        ans = ans[1:-1].strip()
    if ans.endswith(".") and " " not in ans:
        ans = ans[:-1]
    return ans

# ─────────────────────────────────────────────
# HELPERS
# ─────────────────────────────────────────────

def fetch_questions():
    r = requests.get(f"{API_BASE_URL}/questions", timeout=30)
    r.raise_for_status()
    return r.json()


def download_file(task_id: str, file_name: str):
    if not file_name:
        return None
    os.makedirs("./downloads", exist_ok=True)
    path = f"./downloads/{task_id}_{file_name}"
    if os.path.exists(path):
        return path
    try:
        r = requests.get(f"{API_BASE_URL}/files/{task_id}", stream=True, timeout=30)
        if r.status_code == 200:
            with open(path, "wb") as f:
                for chunk in r.iter_content(8192):
                    f.write(chunk)
            return path
    except Exception:
        pass
    return None


def run_task(question: str, file_path: str | None, retries: int = 2) -> str:
    prompt = question
    if file_path:
        prompt += f"\n\nAn evaluation asset file is available locally at: {file_path}"
    prompt += "\n\nOutput only the raw isolated final value. No conversational filler or labels."

    for attempt in range(retries + 1):
        try:
            raw = agent.run(prompt)
            return clean_answer(str(raw))
        except Exception as e:
            print(f"  Attempt {attempt + 1} failed: {e}")
            if attempt < retries:
                time.sleep(5)
    return "ERROR"

# ─────────────────────────────────────────────
# EVALUATION RUNNER
# ─────────────────────────────────────────────

def evaluate(hf_username: str, code_url: str, progress=ui.Progress()):
    if not hf_username.strip() or not code_url.strip():
        return "❌ Please provide your HF username and Space code URL."

    try:
        progress(0, desc="Fetching questions...")
        questions = fetch_questions()
        total = len(questions)
        answers = []

        for i, item in enumerate(questions):
            task_id = item["task_id"]
            question = item["question"]
            file_name = item.get("file_name") or ""

            progress(i / total, desc=f"Task {i+1}/{total}")
            print(f"\n{'='*60}\nTask {i+1}/{total} | {task_id}")

            file_path = download_file(task_id, file_name)
            answer = run_task(question, file_path)

            print(f"A: {repr(answer)}")
            answers.append({"task_id": task_id, "submitted_answer": answer})
            
            # Pacing delay to stay well under free tier limits
            time.sleep(5)

        progress(0.95, desc="Submitting...")
        resp = requests.post(
            f"{API_BASE_URL}/submit",
            json={"username": hf_username.strip(),
                  "agent_code": code_url.strip(),
                  "answers": answers},
            headers={"Content-Type": "application/json"},
            timeout=60,
        )

        if resp.status_code == 200:
            data = resp.json()
            score = data.get("score", 0)
            correct = int(round(score * total))
            return (
                f"### 🎉 Done!\n\n"
                f"**Score:** {score*100:.1f}% ({correct}/{total} Correct)\n\n"
                f"**Server:** {data.get('message', '')}"
            )
        return f"❌ HTTP {resp.status_code}: {resp.text}"

    except Exception as e:
        return f"❌ Error: {e}"

# ─────────────────────────────────────────────
# UI
# ─────────────────────────────────────────────

with ui.Blocks(theme=ui.themes.Soft()) as interface:
    ui.Markdown("# 🏆 GAIA Benchmark Agent — Unit 4")
    ui.Markdown("Runs 20 Level-1 GAIA tasks using web search, Wikipedia, file reading, and an isolated code sandbox.")
    with ui.Row():
        username = ui.Textbox(label="HF Username", placeholder="ramanand-r")
        code_url = ui.Textbox(label="Space Code URL", placeholder="https://huggingface.co/spaces/ramanand-r/Final_Assignment_Template/tree/main")
    btn = ui.Button("🚀 Run & Submit", variant="primary")
    out = ui.Markdown("Enter your details above and click Run.")
    btn.click(fn=evaluate, inputs=[username, code_url], outputs=out)

if __name__ == "__main__":
    interface.launch(debug=True)