import gradio as gr
import requests
import os
import io
import re
import base64
import time
from typing import Optional

try:
    import PyPDF2
    HAS_PYPDF2 = True
except ImportError:
    HAS_PYPDF2 = False

# === CONFIGURATION ===
API_BASE = "https://agents-course-unit4-scoring.hf.space"
HF_API_KEY  = os.getenv("HF_TOKEN")
HF_USERNAME = os.getenv("HF_USERNAME")
GROQ_KEY    = os.getenv("GROQ_API_KEY")
TAVILY_KEY  = os.getenv("TAVILY_API_KEY")

# ─── GROQ ─────────────────────────────────────────────────────────────────────

def call_groq(messages: list, max_tokens: int = 512, system: str = None, temp: float = 0.0) -> str:
    if not GROQ_KEY:
        return ""
    payload = []
    if system:
        payload.append({"role": "system", "content": system})
    payload.extend(messages)

    for model in ["llama-3.3-70b-versatile", "llama-3.1-70b-versatile"]:
        try:
            r = requests.post(
                "https://api.groq.com/openai/v1/chat/completions",
                headers={"Authorization": f"Bearer {GROQ_KEY}", "Content-Type": "application/json"},
                json={"model": model, "messages": payload, "temperature": temp, "max_tokens": max_tokens},
                timeout=60,
            )
            if r.status_code == 200:
                return r.json()["choices"][0]["message"]["content"].strip()
            elif r.status_code == 429:
                time.sleep(5)
        except Exception as e:
            print(f"  Groq error ({model}): {e}")
    return ""

# ─── SEARCH TOOLS ─────────────────────────────────────────────────────────────

def tavily_search(query: str, max_results: int = 5) -> str:
    if not TAVILY_KEY:
        return ""
    try:
        r = requests.post(
            "https://api.tavily.com/search",
            json={"api_key": TAVILY_KEY, "query": query, "max_results": max_results, "search_depth": "advanced"},
            timeout=30,
        )
        if r.status_code == 200:
            items = r.json().get("results", [])
            parts = [f"[{x.get('title','')}]: {x.get('content','')[:600]}" for x in items]
            return "\n\n".join(parts)
    except Exception as e:
        print(f"  Tavily error: {e}")
    return ""


def ddg_search(query: str) -> str:
    try:
        r = requests.get(
            "https://api.duckduckgo.com/",
            params={"q": query, "format": "json", "no_html": 1, "skip_disambig": 1},
            timeout=15, headers={"User-Agent": "Mozilla/5.0"},
        )
        data = r.json()
        parts = []
        if data.get("AbstractText"):
            parts.append(data["AbstractText"])
        for t in data.get("RelatedTopics", [])[:5]:
            if isinstance(t, dict) and t.get("Text"):
                parts.append(t["Text"])
        return "\n\n".join(parts)
    except Exception as e:
        print(f"  DDG error: {e}")
    return ""


def wikipedia_search(query: str) -> str:
    """Search Wikipedia directly for factual questions."""
    try:
        # First find the right page
        r = requests.get(
            "https://en.wikipedia.org/w/api.php",
            params={"action": "query", "list": "search", "srsearch": query,
                    "format": "json", "srlimit": 3},
            timeout=15, headers={"User-Agent": "GaiaAgent/1.0"},
        )
        results = r.json().get("query", {}).get("search", [])
        if not results:
            return ""

        # Get extract of top result
        title = results[0]["title"]
        r2 = requests.get(
            "https://en.wikipedia.org/w/api.php",
            params={"action": "query", "prop": "extracts", "exintro": True,
                    "explaintext": True, "titles": title, "format": "json"},
            timeout=15, headers={"User-Agent": "GaiaAgent/1.0"},
        )
        pages = r2.json().get("query", {}).get("pages", {})
        for page in pages.values():
            extract = page.get("extract", "")
            if extract:
                return f"[Wikipedia — {title}]: {extract[:1500]}"
    except Exception as e:
        print(f"  Wikipedia error: {e}")
    return ""


def smart_search(question: str, query: str) -> str:
    """Run multiple searches and combine results."""
    results = []

    # 1. Tavily (best)
    if TAVILY_KEY:
        t = tavily_search(query, max_results=4)
        if t:
            results.append(t)

    # 2. Wikipedia for factual/historical questions
    wiki_keywords = ["who", "what", "when", "where", "born", "died", "award",
                     "won", "published", "film", "book", "song", "country",
                     "capital", "president", "author", "founded", "invented"]
    if any(w in question.lower() for w in wiki_keywords):
        w = wikipedia_search(query)
        if w:
            results.append(w)

    # 3. DDG fallback
    if not results:
        d = ddg_search(query)
        if d:
            results.append(d)

    return "\n\n---\n\n".join(results)[:4000]

# ─── FILE HANDLING ────────────────────────────────────────────────────────────

def transcribe_audio(audio_bytes: bytes) -> str:
    for model in ["openai/whisper-large-v3", "openai/whisper-base"]:
        try:
            r = requests.post(
                f"https://api-inference.huggingface.co/models/{model}",
                headers={"Authorization": f"Bearer {HF_API_KEY}"},
                data=audio_bytes, timeout=120,
            )
            if r.status_code == 200:
                text = r.json().get("text", "")
                if text:
                    return text
        except Exception as e:
            print(f"  Whisper error: {e}")
    return "Transcription failed."


def extract_pdf(pdf_bytes: bytes) -> str:
    if not HAS_PYPDF2:
        return ""
    try:
        reader = PyPDF2.PdfReader(io.BytesIO(pdf_bytes))
        pages = [p.extract_text() for p in reader.pages if p.extract_text()]
        return "\n\n".join(pages)[:6000]
    except Exception as e:
        return f"PDF error: {e}"


def describe_image_blip(image_bytes: bytes) -> str:
    try:
        r = requests.post(
            "https://api-inference.huggingface.co/models/Salesforce/blip-image-captioning-large",
            headers={"Authorization": f"Bearer {HF_API_KEY}"},
            data=image_bytes, timeout=60,
        )
        if r.status_code == 200:
            result = r.json()
            if isinstance(result, list) and result:
                return result[0].get("generated_text", "")
    except Exception as e:
        print(f"  BLIP error: {e}")
    return ""


def load_task_file(task_id: str, file_name: str, question: str = "") -> str:
    try:
        r = requests.get(f"{API_BASE}/files/{task_id}", timeout=45)
        if r.status_code != 200:
            return ""
        content = r.content
        ct = r.headers.get("content-type", "").lower()
        fn = file_name.lower()

        if "audio" in ct or fn.endswith((".mp3", ".wav", ".flac")):
            print(f"  🎵 Transcribing audio...")
            return f"[AUDIO TRANSCRIPT]\n{transcribe_audio(content)}"
        elif "pdf" in ct or fn.endswith(".pdf"):
            print(f"  📄 Extracting PDF...")
            return f"[PDF CONTENT]\n{extract_pdf(content)}"
        elif "image" in ct or fn.endswith((".jpg", ".jpeg", ".png", ".webp", ".gif")):
            print(f"  🖼️  Describing image...")
            caption = describe_image_blip(content)
            # Also ask Groq to describe it via base64 if supported
            b64 = base64.b64encode(content).decode()
            groq_desc = call_groq([{
                "role": "user",
                "content": f"This is a base64 image. The image captioning model says: '{caption}'. "
                           f"Given the question '{question}', what specific details in the image are relevant? "
                           f"List every visible element, number, text, and object."
            }], max_tokens=300)
            return f"[IMAGE]\nCaption: {caption}\nAnalysis: {groq_desc}"
        else:
            return f"[FILE CONTENT]\n{content.decode('utf-8', errors='replace')[:5000]}"
    except Exception as e:
        print(f"  File error: {e}")
        return ""

# ─── CORE SOLVER ──────────────────────────────────────────────────────────────

SYSTEM_PROMPT = """You are an expert research assistant solving GAIA benchmark questions.
You have access to web search results and file contents.

CRITICAL RULES:
1. Reply with ONLY the final answer — no explanation, no prefix, no "The answer is"
2. Use EXACTLY the format requested in the question
3. For comma-separated lists: word1, word2, word3
4. For numbers: digits only (no units unless asked)
5. For yes/no: Yes or No
6. Base your answer STRICTLY on the provided context — do NOT invent facts
7. If the context clearly contains the answer, extract it precisely
8. Think carefully before answering"""


def build_search_query(question: str) -> str:
    """Ask Groq to write the best search query."""
    resp = call_groq([{
        "role": "user",
        "content": f"Write a 5-7 word web search query to find the answer to:\n{question}\n\nReturn ONLY the search query:"
    }], max_tokens=40)
    resp = re.sub(r'["\'\n]', '', resp).strip()
    if not resp or len(resp) > 150:
        # Fallback: first 7 meaningful words
        words = re.sub(r'[^\w\s]', '', question).split()[:7]
        return " ".join(words)
    return resp


def validate_and_clean(answer: str, question: str) -> str:
    """Ask Groq to validate the answer format matches what the question asks for."""
    if not answer:
        return "UNKNOWN"

    # Strip prefixes
    prefixes = [
        "FINAL ANSWER:", "Final Answer:", "final answer:",
        "Answer:", "answer:", "The answer is", "The answer is:",
        "Based on the context,", "Based on the search results,",
        "According to", "The result is", "It is", "That is",
        "Here is", "Here's", "I believe", "I think",
    ]
    for pfx in prefixes:
        if answer.lower().startswith(pfx.lower()):
            answer = answer[len(pfx):].strip(" :")
            break

    # Remove surrounding quotes
    if len(answer) > 2 and answer[0] == answer[-1] and answer[0] in ('"', "'"):
        answer = answer[1:-1]

    # Take only first meaningful line (avoid leaked reasoning)
    lines = [l.strip() for l in answer.strip().splitlines() if l.strip()]
    if len(lines) > 1:
        # Find the last line that doesn't look like reasoning
        reasoning_starters = ("step", "first", "then", "so ", "therefore",
                               "thus", "note", "let", "we ", "i ", "the question")
        for line in reversed(lines):
            if not any(line.lower().startswith(w) for w in reasoning_starters):
                answer = line
                break

    # Clean trailing period
    if answer.endswith(".") and not re.search(r'\b[A-Z]{1,3}\.$', answer):
        answer = answer[:-1]

    # Normalize list format
    if "comma-separated" in question.lower() or "list" in question.lower():
        answer = re.sub(r'\s+and\s+(?=[^,])', ', ', answer)
        answer = answer.strip(", ")

    return answer.strip()


def solve_question(task_id: str, question: str, file_info: Optional[dict]) -> str:
    print(f"\n{'='*60}")
    print(f"Q: {question[:130]}")

    context_parts = []

    # ── Step 1: Load file ─────────────────────────────────────
    if file_info and file_info.get("file_name"):
        file_ctx = load_task_file(task_id, file_info["file_name"], question)
        if file_ctx:
            context_parts.append(file_ctx)
            print(f"  📁 File: {len(file_ctx)} chars")

    # ── Step 2: Multi-pass web search ─────────────────────────
    query1 = build_search_query(question)
    print(f"  🔍 Search 1: '{query1}'")
    search1 = smart_search(question, query1)
    if search1:
        context_parts.append(f"[SEARCH: '{query1}']\n{search1}")

    # Second search with different angle if first is weak
    if not search1 or len(search1) < 200:
        # Try direct Wikipedia with key entities
        query2 = re.sub(r'\b(what|who|when|where|which|how|the|a|an|is|are|was|were)\b', '',
                        question.lower(), flags=re.I).strip()
        query2 = " ".join(query2.split()[:6])
        if query2 and query2 != query1:
            print(f"  🔍 Search 2 (fallback): '{query2}'")
            search2 = smart_search(question, query2)
            if search2:
                context_parts.append(f"[SEARCH 2: '{query2}']\n{search2}")

    # ── Step 3: Build prompt ──────────────────────────────────
    context = "\n\n" + "="*40 + "\n\n".join(context_parts) if context_parts else "No context available."

    # Detect expected format
    q_low = question.lower()
    fmt = ""
    if "comma-separated" in q_low:
        fmt = "\n⚠️ FORMAT: comma-separated list only (e.g. apple, banana, cherry)"
    elif "how many" in q_low or "how much" in q_low or "what number" in q_low:
        fmt = "\n⚠️ FORMAT: single number only"
    elif "what year" in q_low:
        fmt = "\n⚠️ FORMAT: 4-digit year only"
    elif q_low.strip()[:3] in ("is ", "are", "was", "did", "doe"):
        fmt = "\n⚠️ FORMAT: Yes or No only"
    elif "first name" in q_low:
        fmt = "\n⚠️ FORMAT: first name only"
    elif "last name" in q_low or "surname" in q_low:
        fmt = "\n⚠️ FORMAT: last name only"

    prompt = f"""QUESTION: {question}
{fmt}

CONTEXT:
{context}

Using only the context above, give the precise answer. Think step by step, then output ONLY the final answer:"""

    raw = call_groq([{"role": "user", "content": prompt}], max_tokens=400, system=SYSTEM_PROMPT)
    answer = validate_and_clean(raw, question)

    # ── Step 4: If still empty, try once more with simpler prompt
    if not answer or answer == "UNKNOWN":
        simple = call_groq([{
            "role": "user",
            "content": f"Based on this context, answer the question with ONLY the answer (no explanation):\n\nContext: {context[:2000]}\n\nQuestion: {question}\n\nAnswer:"
        }], max_tokens=100)
        answer = validate_and_clean(simple, question)

    print(f"  → '{answer}'")
    return answer if answer else "UNKNOWN"

# ─── GRADIO ───────────────────────────────────────────────────────────────────

class GaiaAgent:
    def test_single(self) -> str:
        try:
            q = requests.get(f"{API_BASE}/random-question", timeout=15).json()
            has_file = bool(q.get("file_name"))
            answer = solve_question(q["task_id"], q["question"], q if has_file else None)
            return (
                f"🔍 TEST UNE QUESTION\n\n"
                f"Task ID : {q['task_id']}\n"
                f"Question: {q['question']}\n"
                f"Fichier : {q.get('file_name', 'Aucun')}\n\n"
                f"✅ Réponse : '{answer}'"
            )
        except Exception as e:
            return f"❌ Erreur: {e}"

    def run_evaluation(self, progress=gr.Progress()) -> str:
        if not HF_API_KEY or not HF_USERNAME:
            return "❌ ERREUR: Ajoute HF_TOKEN et HF_USERNAME dans Secrets"
        if not GROQ_KEY:
            return "❌ ERREUR: Ajoute GROQ_API_KEY dans Secrets (gratuit sur console.groq.com)"

        progress(0.02, desc="📥 Chargement des questions...")
        try:
            questions = requests.get(f"{API_BASE}/questions", timeout=30).json()
        except Exception as e:
            return f"❌ Impossible de charger les questions: {e}"

        total = len(questions)
        print(f"✅ {total} questions chargées")
        answers = []
        log_lines = [f"✅ {total} questions chargées\n"]

        for idx, q in enumerate(questions):
            task_id = q["task_id"]
            has_file = bool(q.get("file_name"))
            progress((idx + 1) / total, desc=f"Q{idx+1}/{total}...")

            try:
                answer = solve_question(task_id, q["question"], q if has_file else None)
                answers.append({"task_id": task_id, "submitted_answer": answer})
                log_lines.append(f"Q{idx+1}: '{answer[:70]}'")
            except Exception as e:
                print(f"Erreur Q{idx+1}: {e}")
                answers.append({"task_id": task_id, "submitted_answer": "UNKNOWN"})
                log_lines.append(f"Q{idx+1}: ❌ ERREUR — {str(e)[:60]}")

            time.sleep(1)

        progress(0.95, desc="📤 Soumission...")
        code_link = f"https://huggingface.co/spaces/{HF_USERNAME}/gaia-agent/tree/main"

        try:
            resp = requests.post(
                f"{API_BASE}/submit",
                json={"username": HF_USERNAME, "agent_code": code_link, "answers": answers},
                timeout=60,
            )
            if resp.status_code == 200:
                result = resp.json()
                score = result.get("score", 0)
                correct = result.get("correct_count", score)
                pct = (correct / total) * 100 if total else 0
                passed = "🎉 CERTIFICAT OBTENU!" if pct >= 30 else "⚠️ Besoin de 30%+ pour valider"
                return (
                    f"✅ Évaluation terminée!\n\n"
                    f"📊 Score: {correct}/{total} ({pct:.1f}%) — {passed}\n\n"
                    f"🔗 Code: {code_link}\n"
                    f"📈 Leaderboard: https://huggingface.co/spaces/agents-course/Students_leaderboard\n\n"
                    f"--- Détail ---\n" + "\n".join(log_lines)
                )
            else:
                return f"❌ Erreur soumission ({resp.status_code}): {resp.text}\n\n" + "\n".join(log_lines)
        except Exception as e:
            return f"❌ Soumission échouée: {e}\n\n" + "\n".join(log_lines)


def create_interface():
    agent = GaiaAgent()
    status = [
        f"{'✅' if GROQ_KEY else '❌'} GROQ_API_KEY — requis (gratuit sur console.groq.com)",
        f"{'✅' if TAVILY_KEY else '⚠️'} TAVILY_API_KEY — fortement recommandé (gratuit sur app.tavily.com)",
        f"{'✅' if HF_API_KEY else '❌'} HF_TOKEN",
        f"{'✅' if HF_USERNAME else '❌'} HF_USERNAME",
    ]
    with gr.Blocks(title="GAIA Agent", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
# 🤖 GAIA Benchmark Agent

| Secret | Où l'obtenir | Requis? |
|--------|-------------|---------|
| `HF_TOKEN` | hf.co/settings/tokens | ✅ |
| `HF_USERNAME` | Ton pseudo HF | ✅ |
| `GROQ_API_KEY` | [console.groq.com](https://console.groq.com) | ✅ |
| `TAVILY_API_KEY` | [app.tavily.com](https://app.tavily.com) | ⭐ Recommandé |
""")
        gr.Markdown("### Statut\n" + "\n".join(status))
        with gr.Row():
            run_btn = gr.Button("🚀 Évaluation complète", variant="primary", scale=2)
            test_btn = gr.Button("🔍 Tester une question", variant="secondary")
        output = gr.Textbox(label="Résultats", lines=25, show_copy_button=True)
        run_btn.click(fn=agent.run_evaluation, outputs=output, show_progress=True)
        test_btn.click(fn=agent.test_single, outputs=output)
    return demo


if __name__ == "__main__":
    demo = create_interface()
    demo.launch()