import os
import re
import json
import time
import requests
import gradio as gr
import pandas as pd
import wikipediaapi
from pypdf import PdfReader
from youtube_transcript_api import YouTubeTranscriptApi
from smolagents import (
    CodeAgent,
    InferenceClientModel,
    LiteLLMModel,
    WebSearchTool,
    VisitWebpageTool,
    tool,
)

# --- 1. API Constants & Cache ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
CACHE_FILE = "cached_answers.json"


# --- 2. Advanced Multimodal & Research Tools for GAIA ---

@tool
def search_wikipedia(query: str) -> str:
    """
    Fetches the full summary and sections of a Wikipedia page for an exact entity or topic.
    Use this whenever a question mentions Wikipedia, historical facts, discographies, or biographies.
    Args:
        query: Exact title of the Wikipedia article to look up (e.g., "Mercedes Sosa", "Eliud Kipchoge").
    """
    try:
        wiki = wikipediaapi.Wikipedia(user_agent="GAIA_Agent/1.0", language="en")
        page = wiki.page(query)
        if not page.exists():
            return f"No Wikipedia page found for '{query}'. Try searching via web_search first to get the exact title."
        return f"Title: {page.title}\nURL: {page.fullurl}\nContent:\n{page.text[:18000]}"
    except Exception as e:
        return f"Wikipedia error: {e}"


@tool
def get_youtube_transcript(youtube_url_or_id: str) -> str:
    """
    Extracts the full text transcript/captions of a YouTube video.
    Use this whenever a question includes a youtube.com or youtu.be link!
    Args:
        youtube_url_or_id: The full YouTube URL or the 11-character video ID.
    """
    try:
        match = re.search(r"(?:v=|\/)([0-9A-Za-z_-]{11}).*", youtube_url_or_id)
        video_id = match.group(1) if match else youtube_url_or_id.strip()
        ytt_api = YouTubeTranscriptApi()
        fetched = ytt_api.fetch(video_id)
        lines = [f"[{int(snippet.start)}s] {snippet.text}" for snippet in fetched]
        return "\n".join(lines)[:18000]
    except Exception as e:
        return f"Error fetching YouTube transcript: {e}"


@tool
def read_pdf_file(file_path: str) -> str:
    """
    Extracts all text from a local PDF (.pdf) file.
    Args:
        file_path: Path to the local PDF file.
    """
    try:
        reader = PdfReader(file_path)
        text = []
        for idx, page in enumerate(reader.pages):
            text.append(f"--- Page {idx + 1} ---\n{page.extract_text()}")
        return "\n".join(text)[:20000]
    except Exception as e:
        return f"Error reading PDF: {e}"


@tool
def read_text_or_code_file(file_path: str) -> str:
    """
    Reads the text content of a local file (.txt, .py, .csv, .md, .json, .pdb).
    Args:
        file_path: Path to the local file downloaded for the task.
    """
    try:
        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
            return f.read()[:20000]
    except Exception as e:
        return f"Error reading file: {e}"


@tool
def read_excel_or_csv(file_path: str) -> str:
    """
    Reads an Excel (.xlsx, .xls) or CSV (.csv) file and returns its columns, summary stats, and full table data.
    Args:
        file_path: Path to the local Excel or CSV file.
    """
    try:
        if file_path.endswith(".csv"):
            df = pd.read_csv(file_path)
        else:
            df = pd.read_excel(file_path)
        return (
            f"Columns: {list(df.columns)}\n"
            f"Shape: {df.shape}\n"
            f"Data:\n{df.to_string()}\n"
            f"Numeric Summary:\n{df.describe().to_string()}"
        )
    except Exception as e:
        return f"Error reading spreadsheet: {e}"


@tool
def inspect_image_or_audio(file_path: str, task_instruction: str) -> str:
    """
    Uses a multimodal model to transcribe audio (.mp3, .wav) or analyze images (.png, .jpg, .webp) such as chessboards, charts, or handwriting.
    Args:
        file_path: Path to the local audio or image file.
        task_instruction: Specific question or instruction about the image or audio (e.g., "Transcribe this audio word for word" or "List the exact coordinates of all pieces on this chessboard").
    """
    try:
        # Route 1: Gemini Multimodal API (Handles both Audio and Images natively)
        if os.getenv("GEMINI_API_KEY"):
            from google import genai
            client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
            uploaded_file = client.files.upload(file=file_path)
            response = client.models.generate_content(
                model="gemini-2.5-flash",
                contents=[uploaded_file, task_instruction],
            )
            return response.text

        # Route 2: Hugging Face Serverless ASR fallback for audio if no Gemini key
        if file_path.lower().endswith((".mp3", ".wav", ".m4a", ".flac")):
            from huggingface_hub import InferenceClient
            hf_client = InferenceClient(token=os.getenv("HF_TOKEN"))
            result = hf_client.automatic_speech_recognition(file_path, model="openai/whisper-large-v3-turbo")
            return f"Audio Transcription: {result.text}"

        return "Multimodal inspection requires GEMINI_API_KEY secret for images or HF_TOKEN for audio."
    except Exception as e:
        return f"Error inspecting media file: {e}"


# --- 3. Upgraded GAIAAgent with Planning & Two-Stage Answer Verification ---
class GAIAAgent:
    def __init__(self):
        print("Initializing Upgraded GAIAAgent...")

        if os.getenv("GEMINI_API_KEY"):
            print("Using Gemini 2.5 Flash via LiteLLMModel")
            self.model = LiteLLMModel(
                model_id="gemini/gemini-2.5-flash",
                api_key=os.getenv("GEMINI_API_KEY"),
                temperature=0.0,
            )
        else:
            print("Using Qwen/Qwen2.5-Coder-32B-Instruct via InferenceClientModel")
            self.model = InferenceClientModel(
                model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
                token=os.getenv("HF_TOKEN"),
                temperature=0.0,
            )

        self.agent = CodeAgent(
            tools=[
                WebSearchTool(),
                VisitWebpageTool(),
                search_wikipedia,
                get_youtube_transcript,
                read_pdf_file,
                read_text_or_code_file,
                read_excel_or_csv,
                inspect_image_or_audio,
            ],
            model=self.model,
            additional_authorized_imports=[
                "pandas", "numpy", "datetime", "math", "re", "json",
                "csv", "openpyxl", "pypdf", "collections", "itertools",
                "statistics", "urllib", "bs4"
            ],
            max_steps=12,
            planning_interval=3,  # Re-evaluates plan every 3 steps so it never gets stuck
        )

    def clean_and_verify_answer(self, question: str, raw_answer: str) -> str:
        """
        Two-stage post-processor:
        1. Uses an LLM formatter pass to ensure strict GAIA Exact-Match compliance.
        2. Applies regex sanitization as a safety net.
        """
        raw_str = str(raw_answer).strip()

        verification_prompt = f"""You are a strict GAIA benchmark answer extractor.
Given the original QUESTION and the AGENT'S RAW OUTPUT, extract ONLY the exact final answer required by the question.

RULES:
1. Output ONLY the final answer string. Do NOT include any explanation, prefix, or "FINAL ANSWER:".
2. If the question asks for a number/count, output ONLY digits (e.g., 14, not "14 albums" or "fourteen") unless units are explicitly requested.
3. If the question asks for currency/money, follow the unit requested (e.g., if it asks "in USD" or "$", strip extra text).
4. If the question asks for a comma-separated or alphabetical list, double-check the alphabetical/numerical sorting and use ", " (comma + single space) between items.
5. Notice if the question is written backwards/reversed—if so, answer the decoded question directly.
6. Do NOT wrap the output in quotes.

QUESTION:
{question}

AGENT'S RAW OUTPUT:
{raw_str}

EXACT FINAL ANSWER:"""

        try:
            messages = [{"role": "user", "content": [{"type": "text", "text": verification_prompt}]}]
            response = self.model(messages)
            verified = response.content.strip()
        except Exception:
            verified = raw_str

        # Final Regex Safety Net
        ans = re.sub(r"^(final answer|exact final answer)\s*[:\-]?\s*", "", verified, flags=re.IGNORECASE).strip()
        ans = ans.strip('"\'`')
        if ans.endswith(".") and len(ans) < 50:
            ans = ans[:-1].strip()
        return ans

    def __call__(self, question: str, file_path: str | None = None) -> str:
        prompt = f"""You are an elite autonomous agent solving a GAIA benchmark question.

STRATEGY GUIDELINES:
1. Read the question carefully. If the text is reversed or encoded, write a 1-line Python snippet to reverse/decode it first!
2. If a local file path is attached:
   - For `.mp3`, `.wav`, `.png`, `.jpg`: call `inspect_image_or_audio(file_path, task_instruction=...)`.
   - For `.xlsx`, `.csv`: call `read_excel_or_csv(file_path)` or write pandas code directly to filter/sum accurately.
   - For `.pdf`: call `read_pdf_file(file_path)`.
   - For `.py`, `.txt`: call `read_text_or_code_file(file_path)`.
3. If the question mentions a YouTube link, ALWAYS call `get_youtube_transcript` first.
4. If the question asks about Wikipedia, biographies, albums, or historical facts, use `search_wikipedia` in addition to `web_search`.
5. Always use Python code to compute math, counts, sorting, or filtering rather than doing mental math.

QUESTION:
{question}
"""
        if file_path:
            prompt += f"\nATTACHED LOCAL FILE PATH: {file_path}\n"

        prompt += """
When calling `final_answer(...)`, pass ONLY the exact answer required by the question and nothing else.
"""
        try:
            raw_result = self.agent.run(prompt)
            cleaned = self.clean_and_verify_answer(question, raw_result)
            print(f"Raw Output: {raw_result}  -->  ✅ Verified Final Answer: {cleaned}")
            return cleaned
        except Exception as e:
            print(f"❌ Agent execution error: {e}")
            return "Error"


# --- 4. Helper to Download Attached Files via GET /files/{task_id} ---
def download_task_file(api_url: str, task_id: str, file_name: str) -> str | None:
    if not file_name:
        return None
    os.makedirs("downloads", exist_ok=True)
    local_path = os.path.join("downloads", f"{task_id}_{file_name}")
    if os.path.exists(local_path) and os.path.getsize(local_path) > 0:
        return local_path

    file_url = f"{api_url}/files/{task_id}"
    try:
        print(f"Downloading attachment from {file_url} ...")
        resp = requests.get(file_url, timeout=30)
        if resp.status_code == 200:
            with open(local_path, "wb") as f:
                f.write(resp.content)
            return local_path
    except Exception as e:
        print(f"Warning: Failed to download file for task {task_id}: {e}")
    return None


# --- 5. Cache Helpers (Avoid Re-Running Solved Questions) ---
def load_cache() -> dict:
    if os.path.exists(CACHE_FILE):
        try:
            with open(CACHE_FILE, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            return {}
    return {}


def save_cache(cache: dict):
    with open(CACHE_FILE, "w", encoding="utf-8") as f:
        json.dump(cache, f, indent=2, ensure_ascii=False)


# --- 6. Main Evaluation & Submission Loop ---
def run_and_submit_all(profile: gr.OAuthProfile | None, use_cache: bool = True):
    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 above first!", None

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

    try:
        agent = GAIAAgent()
    except Exception as e:
        return f"Error initializing agent: {e}", None

    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main" if space_id else "https://huggingface.co/spaces"

    print(f"Fetching questions from: {questions_url}")
    try:
        response = requests.get(questions_url, timeout=30)
        response.raise_for_status()
        questions_data = response.json()
    except Exception as e:
        return f"Error fetching questions: {e}", None

    cache = load_cache() if use_cache else {}
    results_log = []
    answers_payload = []

    for idx, item in enumerate(questions_data, 1):
        task_id = item.get("task_id")
        question_text = item.get("question")
        file_name = item.get("file_name")

        if not task_id or question_text is None:
            continue

        print(f"\n--- [{idx}/{len(questions_data)}] Task {task_id} ---")

        # Use cached answer if available and valid
        if use_cache and task_id in cache and cache[task_id] not in ("Error", "", None):
            submitted_answer = cache[task_id]
            print(f"⚡ Loaded from cache: {submitted_answer}")
        else:
            local_file = download_task_file(api_url, task_id, file_name)
            submitted_answer = agent(question_text, file_path=local_file)
            cache[task_id] = submitted_answer
            save_cache(cache)
            time.sleep(2)  # Gentle pause to prevent API/search rate-limiting

        answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
        results_log.append({
            "Task ID": task_id,
            "Question": question_text,
            "File": file_name or "None",
            "Submitted Answer": submitted_answer,
        })

    submission_data = {
        "username": username.strip(),
        "agent_code": agent_code,
        "answers": answers_payload,
    }

    print(f"Submitting {len(answers_payload)} answers to {submit_url}...")
    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 Exception as e:
        return f"Submission Failed: {e}", pd.DataFrame(results_log)


# --- 7. Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("# 🕵🏻‍♂️ Upgraded Multimodal GAIA Agent Runner")
    gr.Markdown(
        "1. Log in to your Hugging Face account using the button below.\n"
        "2. Click **Run Evaluation & Submit All Answers** to run the multimodal `CodeAgent` with Wikipedia, YouTube, Audio/Vision, and Two-Stage Exact-Match Verification."
    )
    gr.LoginButton()
    use_cache_checkbox = gr.Checkbox(
        value=True,
        label="Use local answer cache (skips questions already solved in previous runs)"
    )
    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,
        inputs=[gr.State(), use_cache_checkbox],
        outputs=[status_output, results_table],
    )

if __name__ == "__main__":
    demo.launch(debug=True, share=False)