import os
import re
import base64
import tempfile
import requests

from smolagents import CodeAgent, LiteLLMModel, tool, DuckDuckGoSearchTool

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

ANSWER_FORMATTING_RULES = """

## CRITICAL ANSWER FORMATTING RULES (GAIA exact-match grading)

You MUST follow these rules when producing your final answer:

1. Return ONLY the final answer value. No preamble, no explanation, no restated question.
   Never write "The answer is..." or "I think..." — just the raw answer.
2. Never include the literal string "FINAL ANSWER" anywhere in your output.
3. Numbers: Use digits only. No commas, no currency symbols ($, €, £), no percent signs (%),
   no unit abbreviations unless the question explicitly asks for them.
4. Strings: Do NOT include articles (a, an, the) unless they are part of a proper noun.
   Do not abbreviate unless the question explicitly asks for an abbreviation.
5. Comma-separated lists: Apply the above rules to each element, join with ", ".
6. Dates: Use the format explicitly requested. If none specified, use YYYY-MM-DD.
7. For true/false questions, answer only "True" or "False" (no period).
"""


def _clean_answer(raw: str) -> str:
    """Post-process agent output to ensure a clean, gradable answer string."""
    if raw is None:
        return ""
    if not isinstance(raw, str):
        return str(raw)

    answer = raw.strip()

    for prefix in [
        "FINAL ANSWER:", "FINAL ANSWER :", "Final Answer:", "Final Answer :",
        "final answer:", "final answer :",
    ]:
        if answer.startswith(prefix):
            answer = answer[len(prefix):].strip()
            break

    if len(answer) >= 2 and answer[0] == '"' and answer[-1] == '"':
        answer = answer[1:-1].strip()
    if len(answer) >= 2 and answer[0] == "'" and answer[-1] == "'":
        answer = answer[1:-1].strip()

    answer = re.sub(r"\s+", " ", answer).strip()

    return answer


@tool
def download_task_file(task_id: str) -> str:
    """Download the file attached to a GAIA task from the scoring API.

    Args:
        task_id: The task_id string from the questions endpoint.
    """
    url = f"{API_URL}/files/{task_id}"
    try:
        resp = requests.get(url, timeout=60)
        resp.raise_for_status()
    except Exception as e:
        return f"Error downloading file for task {task_id}: {e}"

    content_type = resp.headers.get("content-type", "")
    cd = resp.headers.get("content-disposition", "")
    if "filename=" in cd:
        filename = cd.split("filename=")[-1].strip('" ')
    else:
        ext_map = {
            "pdf": ".pdf", "csv": ".csv", "json": ".json",
            "text": ".txt", "html": ".html", "xml": ".xml",
            "xlsx": ".xlsx", "excel": ".xlsx",
            "png": ".png", "jpeg": ".jpg", "jpg": ".jpg",
            "gif": ".gif", "webp": ".webp",
            "audio": ".mp3", "mpeg": ".mp3", "wav": ".wav",
            "mp4": ".mp4", "ogg": ".ogg", "flac": ".flac",
        }
        ext = ".bin"
        for key, val in ext_map.items():
            if key in content_type:
                ext = val
                break
        filename = f"{task_id}{ext}"

    tmp_dir = tempfile.mkdtemp()
    filepath = os.path.join(tmp_dir, filename)
    with open(filepath, "wb") as f:
        f.write(resp.content)

    return (
        f"File downloaded successfully.\n"
        f"Path: {filepath}\n"
        f"Filename: {filename}\n"
        f"Size: {len(resp.content)} bytes\n"
        f"Content-Type: {content_type}\n\n"
        f"Now use the read_file tool on this path, or analyze_image / transcribe_audio if it is an image or audio file."
    )


@tool
def visit_webpage(url: str) -> str:
    """Visit a webpage and return its text content. Handles common sites like Wikipedia.

    Args:
        url: The full URL of the webpage to visit.
    """
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
        "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
        "Accept-Language": "en-US,en;q=0.9",
    }
    try:
        resp = requests.get(url, headers=headers, timeout=20)
        resp.raise_for_status()
        content_type = resp.headers.get("content-type", "")
        if "text" not in content_type and "html" not in content_type:
            return f"Response is not text (content-type: {content_type}). Cannot read as webpage."
        text = resp.text
        # Simple HTML tag stripping for readability
        import re as _re
        text = _re.sub(r"<script[^>]*>.*?</script>", "", text, flags=_re.DOTALL)
        text = _re.sub(r"<style[^>]*>.*?</style>", "", text, flags=_re.DOTALL)
        text = _re.sub(r"<[^>]+>", " ", text)
        text = _re.sub(r"\s+", " ", text).strip()
        if len(text) > 15000:
            text = text[:15000] + "\n\n[Content truncated at 15000 characters]"
        return text
    except Exception as e:
        return f"Error fetching the webpage: {e}"


@tool
def read_file(file_path: str) -> str:
    """Read and parse a file from disk. Supports PDF, CSV, XLSX, JSON, plain text, and other text formats.

    Args:
        file_path: Absolute path to the file on disk.
    """
    if not os.path.exists(file_path):
        return f"File not found: {file_path}"

    ext = os.path.splitext(file_path)[1].lower()

    try:
        if ext == ".pdf":
            from PyPDF2 import PdfReader
            reader = PdfReader(file_path)
            pages = []
            for i, page in enumerate(reader.pages):
                text = page.extract_text()
                if text:
                    pages.append(f"--- Page {i + 1} ---\n{text}")
            if pages:
                return "\n\n".join(pages)
            return (
                "PDF text extraction returned empty. "
                "This PDF may be image-based. Try converting it to images and using analyze_image."
            )

        if ext == ".csv":
            import pandas as pd
            df = pd.read_csv(file_path)
            return df.to_string(index=False)

        if ext in (".xlsx", ".xls"):
            import pandas as pd
            df = pd.read_excel(file_path)
            return df.to_string(index=False)

        if ext in (".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp", ".tiff"):
            return (
                f"This is an image file at: {file_path}\n"
                "Use the analyze_image tool to examine its contents."
            )

        if ext in (".mp3", ".wav", ".ogg", ".m4a", ".flac", ".mp4", ".aac"):
            return (
                f"This is an audio file at: {file_path}\n"
                "Use the transcribe_audio tool to get its text content."
            )

        if ext in (".json",):
            with open(file_path, "r", encoding="utf-8") as f:
                return f.read()

        if ext in (".zip",):
            import zipfile
            with zipfile.ZipFile(file_path, "r") as z:
                names = z.namelist()
                result = f"ZIP archive contains {len(names)} files:\n"
                for name in names[:50]:
                    result += f"  - {name}\n"
                if len(names) > 50:
                    result += f"  ... and {len(names) - 50} more.\n"
                return result

        with open(file_path, "r", encoding="utf-8", errors="replace") as f:
            return f.read()

    except Exception as e:
        return f"Error reading file {file_path}: {e}"


@tool
def analyze_image(file_path: str, question: str = "Describe this image in detail. Extract all text, numbers, labels, and data you can see.") -> str:
    """Analyze an image using a vision-capable LLM. Use this for photos, charts, diagrams, screenshots, etc.

    Args:
        file_path: Absolute path to the image file.
        question: What you want to know about the image.
    """
    if not os.path.exists(file_path):
        return f"File not found: {file_path}"

    try:
        import litellm

        with open(file_path, "rb") as f:
            img_bytes = f.read()
        img_b64 = base64.b64encode(img_bytes).decode()

        ext = os.path.splitext(file_path)[1].lower()
        mime_map = {
            ".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg",
            ".gif": "image/gif", ".webp": "image/webp", ".bmp": "image/bmp",
            ".tiff": "image/tiff",
        }
        mime = mime_map.get(ext, "image/png")

        model_id = os.getenv("MODEL_ID", "gpt-4o")
        api_key = (
            os.getenv("OPENAI_API_KEY")
            or os.getenv("ANTHROPIC_API_KEY")
            or os.getenv("HF_TOKEN")
        )

        resp = litellm.completion(
            model=model_id,
            api_key=api_key,
            messages=[{
                "role": "user",
                "content": [
                    {"type": "text", "text": question},
                    {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{img_b64}"}},
                ],
            }],
            temperature=0,
        )
        return resp.choices[0].message.content
    except Exception as e:
        return f"Error analyzing image: {e}"


@tool
def transcribe_audio(file_path: str) -> str:
    """Transcribe an audio file to text using OpenAI Whisper API.

    Args:
        file_path: Absolute path to the audio file.
    """
    if not os.path.exists(file_path):
        return f"File not found: {file_path}"

    try:
        from openai import OpenAI

        api_key = os.getenv("OPENAI_API_KEY") or os.getenv("HF_TOKEN")
        client = OpenAI(api_key=api_key) if api_key else OpenAI()

        with open(file_path, "rb") as f:
            result = client.audio.transcriptions.create(
                model="whisper-1",
                file=f,
            )
        return result.text
    except Exception as e:
        return f"Error transcribing audio: {e}"


def build_question_prompt(item: dict) -> str:
    """Build the full prompt to send to the agent, including file hints."""
    question_text = item.get("question", "")
    task_id = item.get("task_id", "")
    file_name = item.get("file_name", "")

    parts = [question_text]
    parts.append(f"\n[TASK_ID: {task_id}]")

    if file_name:
        parts.append(
            f"This task has an attached file: '{file_name}'. "
            f"You MUST first call the download_task_file tool with task_id='{task_id}' "
            f"to download it, then read or analyze the downloaded file before answering."
        )

    return "\n".join(parts)


def create_agent():
    """Create and return a configured CodeAgent for GAIA benchmark questions."""
    model_id = os.getenv("MODEL_ID", "gpt-4o")
    api_key = (
        os.getenv("OPENAI_API_KEY")
        or os.getenv("ANTHROPIC_API_KEY")
        or os.getenv("HF_TOKEN")
    )

    model = LiteLLMModel(
        model_id=model_id,
        api_key=api_key,
        api_base=os.getenv("API_BASE"),
        temperature=0,
    )

    tools = [
        DuckDuckGoSearchTool(),
        visit_webpage,
        download_task_file,
        read_file,
        analyze_image,
        transcribe_audio,
    ]

    agent = CodeAgent(
        tools=tools,
        model=model,
        max_steps=12,
        additional_authorized_imports=[
            "math", "datetime", "json", "re", "os", "collections",
            "itertools", "functools", "string", "random",
            "pandas", "numpy", "requests",
        ],
    )

    agent.prompt_templates["system_prompt"] += ANSWER_FORMATTING_RULES

    return agent