import os
import re
import json
import base64
import tempfile
import subprocess
import mimetypes

import requests
import pandas as pd

from bs4 import BeautifulSoup
from groq import Groq


SCORING_API = "https://agents-course-unit4-scoring.hf.space"

GROQ_API_KEY = os.getenv("GROQ_API_KEY")

VISION_MODEL = os.getenv(
    "GROQ_VISION_MODEL",
    "qwen/qwen3.8-27b"
)

AUDIO_MODEL = os.getenv(
    "GROQ_AUDIO_MODEL",
    "whisper-large-v3-turbo"
)

groq_client = Groq(
    api_key=GROQ_API_KEY
)


# ---------------------------------------------------------
# Web search state
# ---------------------------------------------------------

LAST_SEARCH_RESULTS = []


# ---------------------------------------------------------
# Web search
# ---------------------------------------------------------

def web_search(query):

    global LAST_SEARCH_RESULTS

    try:
        from ddgs import DDGS

        results = []

        with DDGS() as ddgs:
            search_results = ddgs.text(
                query,
                max_results=6
            )

            for item in search_results:

                result = {
                    "title": item.get("title", ""),
                    "url": item.get("href", ""),
                    "snippet": item.get("body", "")
                }

                results.append(result)

        LAST_SEARCH_RESULTS = results

        if not results:
            return "No search results found."

        output = []

        for i, item in enumerate(results):

            output.append(
                f"[{i}]\n"
                f"TITLE: {item['title']}\n"
                f"URL: {item['url']}\n"
                f"SNIPPET: {item['snippet']}"
            )

        return "\n\n".join(output)

    except Exception as e:

        return f"Web search error: {type(e).__name__}: {e}"


# ---------------------------------------------------------
# Read webpage
# ---------------------------------------------------------

def read_webpage(
    url=None,
    cursor=None,
    pattern=None
):

    global LAST_SEARCH_RESULTS

    # -----------------------------------------------------
    # Resolve missing URL using latest search results
    # -----------------------------------------------------

    if not url:

        if cursor is None:
            cursor = 0

        try:
            cursor = int(cursor)
        except Exception:
            cursor = 0

        if (
            cursor < 0
            or cursor >= len(LAST_SEARCH_RESULTS)
        ):
            return (
                "ERROR: No valid URL was supplied and the requested "
                "search-result cursor does not exist."
            )

        url = LAST_SEARCH_RESULTS[cursor].get("url", "")

    if not url:
        return "ERROR: Could not resolve webpage URL."

    # -----------------------------------------------------
    # Fetch webpage
    # -----------------------------------------------------

    try:

        response = requests.get(
            url,
            timeout=20,
            headers={
                "User-Agent": (
                    "Mozilla/5.0 "
                    "(compatible; GAIA-Agent/1.0)"
                )
            }
        )

        response.raise_for_status()

        soup = BeautifulSoup(
            response.text,
            "html.parser"
        )

        for tag in soup(
            [
                "script",
                "style",
                "noscript",
                "svg",
                "nav",
                "footer",
                "header"
            ]
        ):
            tag.decompose()

        text = soup.get_text(
            "\n",
            strip=True
        )

        # -------------------------------------------------
        # If model supplied a pattern, return context
        # around occurrences rather than the whole page.
        # -------------------------------------------------

        if pattern:

            pattern_lower = pattern.lower()
            text_lower = text.lower()

            positions = []

            start = 0

            while True:

                position = text_lower.find(
                    pattern_lower,
                    start
                )

                if position == -1:
                    break

                positions.append(position)

                start = position + len(pattern_lower)

                if len(positions) >= 5:
                    break

            if positions:

                chunks = []

                for position in positions:

                    left = max(
                        0,
                        position - 1800
                    )

                    right = min(
                        len(text),
                        position + 3500
                    )

                    chunks.append(
                        text[left:right]
                    )

                text = "\n\n--- MATCH ---\n\n".join(
                    chunks
                )

        if len(text) > 18000:
            text = text[:18000]

        return (
            f"URL: {url}\n\n"
            f"CONTENT:\n{text}"
        )

    except Exception as e:

        return (
            f"Webpage read error: "
            f"{type(e).__name__}: {e}"
        )


# ---------------------------------------------------------
# Download GAIA attachment
# ---------------------------------------------------------

def download_file(task_id):

    try:

        url = f"{SCORING_API}/files/{task_id}"

        response = requests.get(
            url,
            timeout=60
        )

        response.raise_for_status()

        content_type = response.headers.get(
            "content-type",
            ""
        ).lower()

        disposition = response.headers.get(
            "content-disposition",
            ""
        )

        filename = None

        # -------------------------------------------------
        # Try Content-Disposition
        # -------------------------------------------------

        match = re.search(
            r'filename="?([^";]+)"?',
            disposition,
            re.IGNORECASE
        )

        if match:
            filename = match.group(1)

        # -------------------------------------------------
        # Try URL filename
        # -------------------------------------------------

        if not filename:

            url_name = url.rstrip("/").split("/")[-1]

            if "." in url_name:
                filename = url_name

        # -------------------------------------------------
        # Content-type fallback
        # -------------------------------------------------

        if not filename:

            extension_map = {
                "audio/mpeg": ".mp3",
                "audio/mp3": ".mp3",
                "audio/wav": ".wav",
                "audio/x-wav": ".wav",
                "image/png": ".png",
                "image/jpeg": ".jpg",
                "image/jpg": ".jpg",
                "application/pdf": ".pdf",
                "text/plain": ".txt",
                "text/x-python": ".py",
                "application/json": ".json",
                "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx",
                "application/vnd.ms-excel": ".xls"
            }

            filename = (
                "gaia_attachment"
                + extension_map.get(
                    content_type,
                    ".bin"
                )
            )

        # -------------------------------------------------
        # Save
        # -------------------------------------------------

        safe_filename = os.path.basename(
            filename
        )

        path = os.path.join(
            tempfile.gettempdir(),
            safe_filename
        )

        with open(
            path,
            "wb"
        ) as f:
            f.write(response.content)

        return path

    except Exception as e:

        raise RuntimeError(
            f"Could not download task file: "
            f"{type(e).__name__}: {e}"
        )


# ---------------------------------------------------------
# Inspect text
# ---------------------------------------------------------

def inspect_text(path):

    with open(
        path,
        "r",
        encoding="utf-8",
        errors="replace"
    ) as f:

        text = f.read()

    if len(text) > 20000:
        text = text[:20000]

    return text


# ---------------------------------------------------------
# Inspect Excel
# ---------------------------------------------------------

def inspect_excel(path):

    sheets = pd.read_excel(
        path,
        sheet_name=None
    )

    output = []

    for sheet_name, df in sheets.items():

        output.append(
            f"=== SHEET: {sheet_name} ==="
        )

        output.append(
            f"Rows: {len(df)}"
        )

        output.append(
            f"Columns: {list(df.columns)}"
        )

        output.append(
            df.head(100).to_string(
                index=False
            )
        )

    return "\n\n".join(output)


# ---------------------------------------------------------
# Inspect PDF
# ---------------------------------------------------------

def inspect_pdf(path):

    from PyPDF2 import PdfReader

    reader = PdfReader(path)

    pages = []

    for page in reader.pages:

        text = page.extract_text() or ""

        pages.append(text)

        if len("\n".join(pages)) > 20000:
            break

    text = "\n".join(pages)

    return text[:20000]


# ---------------------------------------------------------
# Inspect Python / source
# ---------------------------------------------------------

def inspect_source(path):

    return inspect_text(path)


# ---------------------------------------------------------
# Generic file inspection
# ---------------------------------------------------------

def inspect_file(task_id):

    path = download_file(task_id)

    extension = (
        os.path.splitext(path)[1]
        .lower()
    )

    if extension in [
        ".xlsx",
        ".xls"
    ]:
        return inspect_excel(path)

    if extension == ".pdf":
        return inspect_pdf(path)

    if extension in [
        ".txt",
        ".py",
        ".json",
        ".csv",
        ".md",
        ".xml",
        ".html"
    ]:
        return inspect_source(path)

    if extension in [
        ".png",
        ".jpg",
        ".jpeg"
    ]:
        return (
            f"Image file downloaded successfully: {path}\n"
            "Use analyze_image for visual analysis."
        )

    if extension in [
        ".mp3",
        ".wav",
        ".m4a",
        ".ogg"
    ]:
        return (
            f"Audio file downloaded successfully: {path}\n"
            "Use transcribe_audio for speech analysis."
        )

    return (
        f"Downloaded file: {path}\n"
        f"Extension: {extension}"
    )


# ---------------------------------------------------------
# Run Python
# ---------------------------------------------------------

def run_python(code):

    path = os.path.join(
        tempfile.gettempdir(),
        "gaia_agent_code.py"
    )

    with open(
        path,
        "w",
        encoding="utf-8"
    ) as f:
        f.write(code)

    try:

        result = subprocess.run(
            [
                "python",
                path
            ],
            capture_output=True,
            text=True,
            timeout=30
        )

        output = (
            result.stdout
            + "\n"
            + result.stderr
        ).strip()

        if len(output) > 12000:
            output = output[:12000]

        return output

    except subprocess.TimeoutExpired:

        return "Python execution timed out after 30 seconds."

    except Exception as e:

        return (
            f"Python execution error: "
            f"{type(e).__name__}: {e}"
        )


# ---------------------------------------------------------
# Audio transcription
# ---------------------------------------------------------

def transcribe_audio(task_id):

    path = download_file(task_id)

    with open(
        path,
        "rb"
    ) as audio_file:

        transcription = groq_client.audio.transcriptions.create(
            file=audio_file,
            model=AUDIO_MODEL,
            response_format="text"
        )

    return str(transcription)


# ---------------------------------------------------------
# Image analysis
# ---------------------------------------------------------

def analyze_image(
    task_id,
    question
):

    path = download_file(task_id)

    extension = (
        os.path.splitext(path)[1]
        .lower()
    )

    mime_map = {
        ".png": "image/png",
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".webp": "image/webp"
    }

    mime_type = mime_map.get(
        extension,
        "image/png"
    )

    with open(
        path,
        "rb"
    ) as f:

        encoded = base64.b64encode(
            f.read()
        ).decode("utf-8")

    response = groq_client.chat.completions.create(

        model=VISION_MODEL,

        temperature=0,

        max_completion_tokens=1500,

        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": (
                            "Analyze this image carefully.\n\n"
                            f"Question: {question}\n\n"
                            "Return the factual answer needed for the question."
                        )
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": (
                                f"data:{mime_type};base64,"
                                f"{encoded}"
                            )
                        }
                    }
                ]
            }
        ]
    )

    return response.choices[0].message.content


# ---------------------------------------------------------
# Tool dispatcher
# ---------------------------------------------------------

def execute_tool(
    name,
    args,
    task_id=None
):

    if name == "web_search":

        return web_search(
            args.get("query", "")
        )

    if name == "read_webpage":

        return read_webpage(
            url=args.get("url"),
            cursor=args.get("cursor"),
            pattern=args.get("pattern")
        )

    if name == "download_file":

        return download_file(
            args.get("task_id") or task_id
        )

    if name == "inspect_file":

        return inspect_file(
            args.get("task_id") or task_id
        )

    if name == "run_python":

        return run_python(
            args.get("code", "")
        )

    if name == "transcribe_audio":

        return transcribe_audio(
            args.get("task_id") or task_id
        )

    if name == "analyze_image":

        return analyze_image(
            task_id=args.get("task_id") or task_id,
            question=args.get("question", "")
        )

    return f"Unknown tool: {name}"