"""Custom tools for the GAIA-style final assignment agent."""
import os

import pandas as pd
import requests
from smolagents import tool

API_URL = "https://agents-course-unit4-scoring.hf.space"
FILES_DIR = "/tmp/task_files"
os.makedirs(FILES_DIR, exist_ok=True)

# Set by app.py once the user is logged in (OAuth token from the Gradio login button).
HF_TOKEN = {"value": None}


@tool
def get_task_file(task_id: str, file_name: str) -> str:
    """Download the file attached to a task and return its local path.

    Args:
        task_id: The task id shown in the question prompt.
        file_name: The attached file name shown in the question prompt.
    """
    path = os.path.join(FILES_DIR, file_name or f"{task_id}.bin")
    if not os.path.exists(path):
        r = requests.get(f"{API_URL}/files/{task_id}", timeout=60)
        if r.status_code != 200:
            return f"FILE UNAVAILABLE (HTTP {r.status_code}). Answer from your own knowledge/web search if possible, else give your best guess."
        with open(path, "wb") as f:
            f.write(r.content)
    return path


@tool
def read_text_file(path: str) -> str:
    """Read a local text/code/csv file and return its contents (truncated at 20000 chars).

    Args:
        path: Local path returned by get_task_file.
    """
    with open(path, "r", errors="replace") as f:
        return f.read()[:20000]


@tool
def read_excel_file(path: str) -> str:
    """Dump every sheet of an .xlsx/.xls/.csv file as CSV text so it can be analysed.

    Args:
        path: Local path returned by get_task_file.
    """
    if path.lower().endswith(".csv"):
        return pd.read_csv(path).to_csv(index=False)[:20000]
    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())[:20000]


@tool
def transcribe_audio(path: str) -> str:
    """Transcribe an audio file (mp3/wav/m4a) to text with Whisper.

    Args:
        path: Local path returned by get_task_file.
    """
    from huggingface_hub import InferenceClient

    client = InferenceClient(token=HF_TOKEN["value"])
    return client.automatic_speech_recognition(path, model="openai/whisper-large-v3").text


@tool
def describe_image(path: str, question: str) -> str:
    """Ask a vision-language model a question about a local image.

    Args:
        path: Local image path returned by get_task_file.
        question: What to find out about the image.
    """
    import base64

    from huggingface_hub import InferenceClient

    with open(path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()
    client = InferenceClient(token=HF_TOKEN["value"])
    resp = client.chat_completion(
        model="Qwen/Qwen2.5-VL-72B-Instruct",
        messages=[{
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}},
                {"type": "text", "text": question},
            ],
        }],
        max_tokens=800,
    )
    return resp.choices[0].message.content


@tool
def youtube_transcript(url: str) -> str:
    """Get the caption transcript of a YouTube video (if captions exist).

    Args:
        url: Full YouTube URL.
    """
    from youtube_transcript_api import YouTubeTranscriptApi

    vid = url.split("v=")[-1].split("&")[0] if "v=" in url else url.rsplit("/", 1)[-1]
    try:
        segs = YouTubeTranscriptApi().fetch(vid)
        return " ".join(s.text for s in segs)[:15000]
    except Exception as e:  # noqa: BLE001
        return f"No transcript available: {e}"


@tool
def wikipedia_page(title: str) -> str:
    """Fetch the plain-text content of an English Wikipedia page by title.

    Args:
        title: Exact page title, e.g. 'Mercedes Sosa'.
    """
    r = requests.get(
        "https://en.wikipedia.org/w/api.php",
        params={"action": "query", "prop": "extracts", "explaintext": 1, "titles": title,
                "format": "json", "redirects": 1},
        headers={"User-Agent": "agents-course-final/1.0"},
        timeout=30,
    )
    try:
        pages = r.json()["query"]["pages"]
        return next(iter(pages.values())).get("extract", "NOT FOUND")[:25000]
    except Exception as e:  # noqa: BLE001
        return f"Wikipedia lookup failed (HTTP {r.status_code}): {e}. Try web search instead."


ALL_TOOLS = [
    get_task_file, read_text_file, read_excel_file,
    transcribe_audio, describe_image, youtube_transcript, wikipedia_page,
]