import base64
import mimetypes
import os
import random
import re
import tempfile
import time
from urllib.parse import parse_qs, unquote, urlparse

import gradio as gr
import pandas as pd
import requests
import spaces
from bs4 import BeautifulSoup
from huggingface_hub import InferenceClient
from openai import OpenAI
from youtube_transcript_api import YouTubeTranscriptApi


DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
ATTACHMENT_MIRROR = (
    "https://huggingface.co/spaces/bstraehle/gaia/resolve/"
    "617431ad4021eb2a62ab81488d1ec1cd7d1e3e19/data"
)
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"
TEXT_MODELS = ("openrouter/free",)
VISION_MODEL = "openrouter/free"
USER_AGENT = "Mozilla/5.0 (compatible; HF-Agents-Course-Student/1.0)"


@spaces.GPU
def zero_gpu_placeholder():
    """Register ZeroGPU support. Agent inference itself uses OpenRouter."""
    return None


def clean_answer(answer: str) -> str:
    """Remove wrappers that would break the benchmark's exact-match scoring."""
    text = (answer or "").strip()
    final_tags = re.findall(r"(?is)<final>\s*(.*?)\s*</final>", text)
    if final_tags:
        text = final_tags[-1]
    else:
        labelled_lines = re.findall(
            r"(?im)^\s*(?:final\s+answer|answer)\s*:\s*(.+?)\s*$",
            text,
        )
        if labelled_lines:
            text = labelled_lines[-1]
    text = re.sub(r"(?i)^\s*(final\s+answer|answer)\s*:\s*", "", text)
    text = text.replace("FINAL ANSWER", "").strip()
    text = text.strip("`\"'")
    if text.startswith("[") and text.endswith("]"):
        text = text[1:-1].strip()
    return " ".join(text.splitlines()).strip()


def safe_public_url(url: str) -> bool:
    """Reject local/private destinations before fetching model-discovered URLs."""
    try:
        parsed = urlparse(url)
        host = (parsed.hostname or "").lower()
        return (
            parsed.scheme in {"http", "https"}
            and host not in {"localhost", "127.0.0.1", "0.0.0.0", "::1"}
            and not host.endswith(".local")
            and not re.match(r"^(10\.|192\.168\.|172\.(1[6-9]|2\d|3[01])\.)", host)
        )
    except Exception:
        return False


def normalize_ddg_url(url: str) -> str:
    if "duckduckgo.com/l/" in url:
        target = parse_qs(urlparse(url).query).get("uddg", [url])[0]
        return unquote(target)
    return url


def web_search(query: str, limit: int = 5) -> list[dict]:
    """Search DuckDuckGo HTML, with Wikipedia API as a reliable fallback."""
    results = []
    try:
        response = requests.post(
            "https://html.duckduckgo.com/html/",
            data={"q": query[:500]},
            headers={"User-Agent": USER_AGENT},
            timeout=15,
        )
        response.raise_for_status()
        soup = BeautifulSoup(response.text, "html.parser")
        for item in soup.select(".result"):
            anchor = item.select_one("a.result__a")
            if not anchor:
                continue
            url = normalize_ddg_url(anchor.get("href", ""))
            if not safe_public_url(url):
                continue
            snippet_node = item.select_one(".result__snippet")
            results.append(
                {
                    "title": anchor.get_text(" ", strip=True),
                    "url": url,
                    "snippet": snippet_node.get_text(" ", strip=True) if snippet_node else "",
                }
            )
            if len(results) >= limit:
                break
    except Exception as exc:
        print(f"DuckDuckGo search failed: {exc}")

    if results:
        return results

    try:
        response = requests.get(
            "https://en.wikipedia.org/w/api.php",
            params={
                "action": "query",
                "list": "search",
                "srsearch": query[:300],
                "format": "json",
                "utf8": 1,
                "srlimit": limit,
            },
            headers={"User-Agent": USER_AGENT},
            timeout=15,
        )
        response.raise_for_status()
        for item in response.json().get("query", {}).get("search", []):
            title = item["title"]
            results.append(
                {
                    "title": title,
                    "url": f"https://en.wikipedia.org/wiki/{title.replace(' ', '_')}",
                    "snippet": BeautifulSoup(item.get("snippet", ""), "html.parser").get_text(" "),
                }
            )
    except Exception as exc:
        print(f"Wikipedia search failed: {exc}")
    return results


def fetch_page(url: str, max_chars: int = 6000) -> str:
    if not safe_public_url(url):
        return ""
    try:
        response = requests.get(
            url,
            headers={"User-Agent": USER_AGENT},
            timeout=15,
            allow_redirects=True,
        )
        response.raise_for_status()
        if not safe_public_url(response.url):
            return ""
        content_type = response.headers.get("content-type", "")
        if "html" not in content_type and "text" not in content_type:
            return ""
        soup = BeautifulSoup(response.text, "html.parser")
        for node in soup(["script", "style", "nav", "footer", "form", "noscript"]):
            node.decompose()
        return re.sub(r"\s+", " ", soup.get_text(" ", strip=True))[:max_chars]
    except Exception as exc:
        print(f"Page fetch failed for {url}: {exc}")
        return ""


def youtube_transcript(question: str) -> str:
    match = re.search(r"(?:v=|youtu\.be/)([A-Za-z0-9_-]{11})", question)
    if not match:
        return ""
    video_id = match.group(1)
    try:
        transcript = YouTubeTranscriptApi().fetch(video_id)
        return " ".join(snippet.text for snippet in transcript)[:18000]
    except Exception as exc:
        print(f"YouTube transcript failed: {exc}")
        return ""


def download_attachment(task_id: str, file_name: str) -> str:
    suffix = os.path.splitext(file_name)[1]
    response = requests.get(
        f"{DEFAULT_API_URL}/files/{task_id}",
        headers={"User-Agent": USER_AGENT},
        timeout=30,
    )
    if response.status_code == 404:
        print(f"Course file endpoint is missing {task_id}; using pinned mirror.")
        response = requests.get(
            f"{ATTACHMENT_MIRROR}/{file_name}",
            headers={"User-Agent": USER_AGENT},
            timeout=30,
        )
    response.raise_for_status()
    handle = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
    handle.write(response.content)
    handle.close()
    return handle.name


def describe_image(path: str, question: str, token: str) -> str:
    mime = mimetypes.guess_type(path)[0] or "image/png"
    with open(path, "rb") as image_file:
        data_url = f"data:{mime};base64,{base64.b64encode(image_file.read()).decode()}"
    client = OpenAI(
        base_url=OPENROUTER_BASE_URL,
        api_key=token,
        timeout=120,
        default_headers={
            "HTTP-Referer": "https://huggingface.co/spaces/AbelCtw/Final_Assignment_Template",
            "X-Title": "HF Agents Course Final Assignment",
        },
    )
    response = client.chat.completions.create(
        model=VISION_MODEL,
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": (
                            "Analyze this image carefully for the benchmark question below. "
                            "State all visual facts and then give the most likely exact answer.\n\n"
                            f"Question: {question}"
                        ),
                    },
                    {"type": "image_url", "image_url": {"url": data_url}},
                ],
            }
        ],
        max_tokens=900,
        temperature=0.1,
    )
    return response.choices[0].message.content or ""


def extract_attachment(path: str, file_name: str, question: str, token: str) -> str:
    extension = os.path.splitext(file_name)[1].lower()
    if extension in {".txt", ".py", ".md", ".csv", ".json"}:
        with open(path, "r", encoding="utf-8", errors="replace") as file:
            return file.read()[:24000]
    if extension in {".xlsx", ".xls"}:
        sheets = pd.read_excel(path, sheet_name=None)
        rendered = []
        for sheet_name, frame in sheets.items():
            rendered.append(f"SHEET: {sheet_name}\n{frame.head(500).to_csv(index=False)}")
        return "\n\n".join(rendered)[:30000]
    if extension in {".mp3", ".wav", ".flac", ".m4a", ".ogg"}:
        hf_token = os.getenv("HF_TOKEN")
        if not hf_token:
            raise RuntimeError(
                "Audio transcription needs the existing optional HF_TOKEN secret."
            )
        client = InferenceClient(provider="hf-inference", api_key=hf_token)
        transcription = client.automatic_speech_recognition(
            path,
            model="openai/whisper-large-v3-turbo",
        )
        return getattr(transcription, "text", str(transcription))[:24000]
    if extension in {".png", ".jpg", ".jpeg", ".webp"}:
        return describe_image(path, question, token)
    return f"Attachment {file_name} could not be parsed automatically."


def should_search(question: str, has_attachment: bool) -> bool:
    if has_attachment:
        return False
    reversed_question = question[::-1].lower()
    if "understand this sentence" in reversed_question:
        return False
    if "given this table" in question.lower() or "grocery list" in question.lower():
        return False
    return True


class ResearchAgent:
    def __init__(self, token: str):
        self.token = token
        self.client = OpenAI(
            base_url=OPENROUTER_BASE_URL,
            api_key=token,
            timeout=120,
            default_headers={
                "HTTP-Referer": "https://huggingface.co/spaces/AbelCtw/Final_Assignment_Template",
                "X-Title": "HF Agents Course Final Assignment",
            },
        )

    def ask_model(self, prompt: str) -> str:
        last_error = None
        for attempt in range(4):
            for model in TEXT_MODELS:
                try:
                    response = self.client.chat.completions.create(
                        model=model,
                        messages=[
                            {
                                "role": "system",
                            "content": (
                                "You solve GAIA-style factual questions using supplied evidence. "
                                "Obey the requested format exactly. You may reason briefly, but the "
                                "last line must be ANSWER: followed by only the exact answer, with no "
                                "citation or markdown."
                                ),
                            },
                            {"role": "user", "content": prompt},
                        ],
                        temperature=0.05,
                        max_tokens=500,
                    )
                    answer = clean_answer(response.choices[0].message.content or "")
                    if not answer:
                        raise RuntimeError("OpenRouter returned an empty answer")
                    return answer
                except Exception as exc:
                    last_error = exc
                    print(f"Model {model}, attempt {attempt + 1} failed: {exc}")
            if attempt < 3:
                time.sleep((2**attempt) + random.random())
        raise RuntimeError(f"All text models failed: {last_error}")

    def solve(self, question: str, task_id: str, file_name: str) -> str:
        evidence = []
        temp_path = None
        try:
            if file_name:
                temp_path = download_attachment(task_id, file_name)
                attachment_text = extract_attachment(temp_path, file_name, question, self.token)
                evidence.append(f"ATTACHMENT {file_name}:\n{attachment_text}")

            transcript = youtube_transcript(question)
            if transcript:
                evidence.append(f"YOUTUBE TRANSCRIPT:\n{transcript}")

            if should_search(question, bool(file_name)):
                search_results = web_search(question)
                for index, result in enumerate(search_results, start=1):
                    page_text = fetch_page(result["url"]) if index <= 3 else ""
                    evidence.append(
                        f"WEB RESULT {index}: {result['title']}\n"
                        f"URL: {result['url']}\n"
                        f"SNIPPET: {result['snippet']}\n"
                        f"PAGE TEXT: {page_text}"
                    )

            prompt = (
                f"QUESTION:\n{question}\n\n"
                "EVIDENCE (may contain irrelevant or conflicting text; verify internally):\n"
                + ("\n\n".join(evidence) if evidence else "No external evidence was needed.")
                + "\n\nReturn only the exact answer in the format requested by the question."
            )
            return self.ask_model(prompt[:52000])
        finally:
            if temp_path and os.path.exists(temp_path):
                os.unlink(temp_path)


def run_and_submit_all(
    profile: gr.OAuthProfile | None,
    oauth_token: gr.OAuthToken | None,
):
    """Fetch all questions, solve them, and submit exact-match answers."""
    if profile is None:
        return (
            "Please sign in with Hugging Face first. No answers were submitted.",
            None,
        )

    username = profile.username
    space_id = os.getenv("SPACE_ID")
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
    provider_token = os.getenv("OPENROUTER_API_KEY")
    if not provider_token:
        return (
            "Missing OpenRouter credential. Add a private Space secret named "
            "OPENROUTER_API_KEY. No answers were submitted.",
            None,
        )

    agent = ResearchAgent(provider_token)
    try:
        access_check = agent.ask_model("Reply with exactly: OK")
        if not access_check.strip():
            raise RuntimeError(f"Unexpected preflight response: {access_check!r}")
    except Exception as exc:
        return (
            "Inference preflight failed, so evaluation was stopped before any submission. "
            "Create an OpenRouter API key and save it as the OPENROUTER_API_KEY "
            f"Space secret. Detail: {exc}",
            None,
        )

    try:
        response = requests.get(f"{DEFAULT_API_URL}/questions", timeout=20)
        response.raise_for_status()
        questions = response.json()
    except Exception as exc:
        return f"Could not fetch questions: {exc}", None

    results_log = []
    answers_payload = []
    for index, item in enumerate(questions, start=1):
        task_id = item.get("task_id")
        question = item.get("question", "")
        file_name = item.get("file_name", "")
        print(f"Solving {index}/{len(questions)}: {task_id}")
        try:
            answer = agent.solve(question, task_id, file_name)
            if not answer.strip():
                raise RuntimeError("The model returned an empty answer")
        except Exception as exc:
            print(f"Agent error on {task_id}: {exc}")
            results_log.append(
                {
                    "Task ID": task_id,
                    "Question": question,
                    "Submitted Answer": f"ERROR (not submitted): {exc}",
                }
            )
            return (
                f"Stopped safely at question {index}/{len(questions)}. "
                "No leaderboard submission was made because at least one answer failed. "
                f"Task {task_id}: {exc}",
                pd.DataFrame(results_log),
            )
        answers_payload.append({"task_id": task_id, "submitted_answer": answer})
        results_log.append(
            {
                "Task ID": task_id,
                "Question": question,
                "Submitted Answer": answer,
            }
        )

    if len(answers_payload) != len(questions) or any(
        not item["submitted_answer"].strip() for item in answers_payload
    ):
        return (
            "Safety check stopped submission because the answer set was incomplete.",
            pd.DataFrame(results_log),
        )

    submission = {
        "username": username.strip(),
        "agent_code": agent_code,
        "answers": answers_payload,
    }
    try:
        response = requests.post(
            f"{DEFAULT_API_URL}/submit",
            json=submission,
            timeout=90,
        )
        response.raise_for_status()
        result = response.json()
        status = (
            "Submission Successful!\n"
            f"User: {result.get('username')}\n"
            f"Overall Score: {result.get('score', 'N/A')}% "
            f"({result.get('correct_count', '?')}/{result.get('total_attempted', '?')} correct)\n"
            f"Message: {result.get('message', 'No message received.')}"
        )
    except Exception as exc:
        status = f"Submission failed after answers were generated: {exc}"
    return status, pd.DataFrame(results_log)


with gr.Blocks() as demo:
    gr.Markdown("# Research Agent Evaluation Runner")
    gr.Markdown(
        """
This agent searches the web, reads benchmark attachments, transcribes audio,
and uses OpenRouter's free model router to produce exact-match answers.

1. Add a private Space secret named `OPENROUTER_API_KEY`.
2. Sign in with Hugging Face.
3. Click **Run Evaluation & Submit All Answers** once.
4. Keep this tab open while the 20 questions are processed.
"""
    )
    gr.LoginButton()
    run_button = gr.Button("Run Evaluation & Submit All Answers", variant="primary")
    status_output = gr.Textbox(
        label="Run Status / Submission Result",
        lines=6,
        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__":
    demo.queue().launch(debug=True, share=False)