import os
import re
import sys
import base64
import tempfile
import subprocess
import threading
import concurrent.futures as cf

import gradio as gr
import requests
import pandas as pd
import anthropic
from bs4 import BeautifulSoup

# --- Constants ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
MODEL = os.getenv("CLAUDE_MODEL", "claude-sonnet-5-5")
MAX_STEPS = 14          # tool-use rounds per question
MAX_WORKERS = 3         # questions answered in parallel (lower it if you hit rate limits)
CHUNK = 12000           # characters returned per page fetch

IMAGE_TYPES = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg",
               ".gif": "image/gif", ".webp": "image/webp"}
AUDIO_EXT = (".mp3", ".wav", ".m4a", ".flac", ".ogg")

SYSTEM = """You are a careful research agent solving questions from the GAIA benchmark.
Answers are graded by exact match, so the format matters as much as the content.

Method:
- Think step by step and use the tools whenever a fact is not certain. Do not guess from memory
  when a lookup is possible. For counting, sorting and arithmetic, use run_python.
- Some questions come with an attached file; its local path is given in the message.
  Use read_file (text, Excel, audio transcription) or run_python on it. Images are attached directly.
- Wikipedia pages: fetch_url works well. For tables, run_python with pandas.read_html is more reliable.
- Respect every detail of the question (date ranges, "inclusive", "only first name", alphabetical order, etc.).
- If a question is a puzzle (reversed text, logic table), solve it directly and then apply the instruction it contains.

Output rules:
- End your reply with one line: FINAL ANSWER: <answer>
- The answer must be only the requested value: a number, a short phrase, or a comma-separated list.
- No units, currency symbols or thousands separators unless the question asks for them.
- No articles or abbreviations in strings unless asked. Write numbers as digits.
- If you truly cannot find the answer after trying, still give your best single guess."""

# --- Tools ---------------------------------------------------------------

TOOLS = [
    {"type": "web_search_20250305", "name": "web_search", "max_uses": 6},
    {
        "name": "fetch_url",
        "description": "Download a web page and return its text content (up to ~12000 characters per call). "
                       "Use `start` to read further into long pages.",
        "input_schema": {
            "type": "object",
            "properties": {"url": {"type": "string"}, "start": {"type": "integer", "description": "Character offset, default 0"}},
            "required": ["url"],
        },
    },
    {
        "name": "run_python",
        "description": "Run Python 3 code (pandas, requests, bs4 available) and return stdout/stderr. "
                       "Always print() what you need. 60 second limit.",
        "input_schema": {"type": "object", "properties": {"code": {"type": "string"}}, "required": ["code"]},
    },
    {
        "name": "read_file",
        "description": "Read a local file: text/code/csv/json, Excel (all sheets as CSV), or audio (transcribed to text).",
        "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]},
    },
    {
        "name": "youtube_transcript",
        "description": "Get the transcript (subtitles) of a YouTube video. Cannot see the images of the video.",
        "input_schema": {"type": "object", "properties": {"url": {"type": "string"}}, "required": ["url"]},
    },
]

_whisper = None
_whisper_lock = threading.Lock()


def transcribe(path: str) -> str:
    global _whisper
    with _whisper_lock:
        if _whisper is None:
            from faster_whisper import WhisperModel
            _whisper = WhisperModel("base", device="cpu", compute_type="int8")
        segments, _ = _whisper.transcribe(path)
        return " ".join(s.text.strip() for s in segments)


def tool_fetch_url(url: str, start: int = 0) -> str:
    r = requests.get(url, headers={"User-Agent": "Mozilla/5.0 (compatible; AgentCourseBot/1.0)"}, timeout=25)
    r.raise_for_status()
    if "html" in r.headers.get("content-type", ""):
        soup = BeautifulSoup(r.text, "html.parser")
        for t in soup(["script", "style", "nav", "footer"]):
            t.decompose()
        text = re.sub(r"\n{2,}", "\n", soup.get_text("\n"))
    else:
        text = r.text
    return text[start:start + CHUNK] or "(no more content)"


def tool_run_python(code: str) -> str:
    with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f:
        f.write(code)
    env = {k: v for k, v in os.environ.items() if k not in ("ANTHROPIC_API_KEY", "HF_TOKEN")}
    try:
        p = subprocess.run([sys.executable, f.name], capture_output=True, text=True, timeout=60, env=env)
    except subprocess.TimeoutExpired:
        return "Timed out after 60 seconds."
    out = p.stdout + (f"\nSTDERR:\n{p.stderr}" if p.stderr else "")
    return out[:8000] or "(no output)"


def tool_read_file(path: str) -> str:
    ext = os.path.splitext(path)[1].lower()
    if ext in AUDIO_EXT:
        return transcribe(path)
    if ext in (".xlsx", ".xls"):
        sheets = pd.read_excel(path, sheet_name=None)
        return "\n\n".join(f"## Sheet: {n}\n{df.to_csv(index=False)}" for n, df in sheets.items())[:15000]
    with open(path, "r", errors="replace") as f:
        return f.read()[:15000]


def tool_youtube_transcript(url: str) -> str:
    from youtube_transcript_api import YouTubeTranscriptApi
    m = re.search(r"(?:v=|youtu\.be/)([\w-]{11})", url)
    if not m:
        return "Could not find a video id in that URL."
    data = YouTubeTranscriptApi().fetch(m.group(1))
    return "\n".join(f"[{int(s.start)}s] {s.text}" for s in data)[:15000]


def run_tool(name: str, args: dict) -> str:
    try:
        if name == "fetch_url":
            return tool_fetch_url(**args)
        if name == "run_python":
            return tool_run_python(**args)
        if name == "read_file":
            return tool_read_file(**args)
        if name == "youtube_transcript":
            return tool_youtube_transcript(**args)
        return f"Unknown tool {name}"
    except Exception as e:
        return f"Tool error: {type(e).__name__}: {e}"


# --- Agent ---------------------------------------------------------------

def extract_final(resp) -> str:
    text = "".join(b.text for b in resp.content if b.type == "text").strip()
    matches = re.findall(r"FINAL ANSWER:\s*(.+)", text)
    answer = matches[-1].strip() if matches else text.splitlines()[-1].strip() if text else ""
    return answer.strip().strip("*").rstrip(".") if answer else ""


class ClaudeAgent:
    def __init__(self):
        self.client = anthropic.Anthropic(max_retries=6)
        print(f"ClaudeAgent initialized with model {MODEL}")

    def __call__(self, question: str, file_path: str | None = None) -> str:
        content = []
        text = question
        if file_path:
            ext = os.path.splitext(file_path)[1].lower()
            if ext in IMAGE_TYPES:
                with open(file_path, "rb") as f:
                    data = base64.standard_b64encode(f.read()).decode()
                content.append({"type": "image", "source": {"type": "base64", "media_type": IMAGE_TYPES[ext], "data": data}})
                text += "\n\n(The attached image is shown above.)"
            else:
                text += f"\n\n[Attached file saved locally at: {file_path}]"
        content.append({"type": "text", "text": text})
        messages = [{"role": "user", "content": content}]

        resp = None
        for _ in range(MAX_STEPS):
            resp = self.client.messages.create(
                model=MODEL, max_tokens=3000, system=SYSTEM, tools=TOOLS, messages=messages
            )
            messages.append({"role": "assistant", "content": resp.content})
            if resp.stop_reason == "tool_use":
                results = [
                    {"type": "tool_result", "tool_use_id": b.id, "content": run_tool(b.name, b.input)}
                    for b in resp.content if b.type == "tool_use"
                ]
                messages.append({"role": "user", "content": results})
            elif resp.stop_reason == "pause_turn":
                continue
            else:
                break

        # If we ran out of steps without a final answer, force one.
        full = "".join(b.text for b in resp.content if b.type == "text") if resp else ""
        if "FINAL ANSWER:" not in full:
            messages.append({"role": "user", "content": "Stop using tools. Give your best final answer now, ending with 'FINAL ANSWER: <answer>'."})
            resp = self.client.messages.create(model=MODEL, max_tokens=500, system=SYSTEM, messages=messages)
        return extract_final(resp)


def download_task_file(task_id: str, file_name: str) -> str | None:
    if not file_name:
        return None
    try:
        r = requests.get(f"{DEFAULT_API_URL}/files/{task_id}", timeout=30)
        if r.status_code != 200 or not r.content:
            return None
        path = os.path.join(tempfile.gettempdir(), f"{task_id}_{file_name}")
        with open(path, "wb") as f:
            f.write(r.content)
        return path
    except Exception as e:
        print(f"Could not download file for {task_id}: {e}")
        return None


def run_and_submit_all(profile: gr.OAuthProfile | None):
    space_id = os.getenv("SPACE_ID")

    if profile:
        username = f"{profile.username}"
        print(f"User logged in: {username}")
    else:
        return "Please Login to Hugging Face with the button.", None

    if not os.getenv("ANTHROPIC_API_KEY"):
        return "Missing ANTHROPIC_API_KEY. Add it as a secret in the Space settings.", None

    questions_url = f"{DEFAULT_API_URL}/questions"
    submit_url = f"{DEFAULT_API_URL}/submit"

    try:
        agent = ClaudeAgent()
    except Exception as e:
        return f"Error initializing agent: {e}", None
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"

    try:
        response = requests.get(questions_url, timeout=15)
        response.raise_for_status()
        questions_data = response.json()
        if not questions_data:
            return "Fetched questions list is empty or invalid format.", None
    except Exception as e:
        return f"Error fetching questions: {e}", None

    def solve(item):
        task_id = item.get("task_id")
        question_text = item.get("question")
        if not task_id or question_text is None:
            return None
        try:
            path = download_task_file(task_id, item.get("file_name"))
            answer = agent(question_text, path)
        except Exception as e:
            print(f"Error on task {task_id}: {e}")
            answer = f"AGENT ERROR: {e}"
        print(f"[{task_id}] -> {answer}")
        return task_id, question_text, answer

    with cf.ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
        outputs = [o for o in pool.map(solve, questions_data) if o]

    results_log = [{"Task ID": t, "Question": q, "Submitted Answer": a} for t, q, a in outputs]
    answers_payload = [{"task_id": t, "submitted_answer": a} for t, q, a in outputs if not a.startswith("AGENT ERROR")]

    if not answers_payload:
        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)

    submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
    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.')}"
        )
        return final_status, pd.DataFrame(results_log)
    except requests.exceptions.HTTPError as e:
        detail = f"Server responded with status {e.response.status_code}. Response: {e.response.text[:500]}"
        return f"Submission Failed: {detail}", pd.DataFrame(results_log)
    except Exception as e:
        return f"Submission Failed: {e}", pd.DataFrame(results_log)


# --- Gradio interface ---
with gr.Blocks() as demo:
    gr.Markdown("# Basic Agent Evaluation Runner")
    gr.Markdown(
        """
        **Instructions:**

        1. Add your `ANTHROPIC_API_KEY` as a secret in the Space settings.
        2. Log in to your Hugging Face account using the button below.
        3. Click 'Run Evaluation & Submit All Answers'. It can take 10 to 20 minutes.
        """
    )
    gr.LoginButton()
    run_button = gr.Button("Run Evaluation & Submit All Answers")
    status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, 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__":
    print("Launching Gradio Interface for Agent Evaluation...")
    demo.launch(debug=True, share=False)