import os
import gradio as gr
import requests
import pandas as pd
import google.generativeai as genai
from typing import Optional

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

class ImprovedGAIAgent:
    """Agente melhorado para o GAIA com Chain of Thought."""
    
    def __init__(self):
        print("🚀 Inicializando agente melhorado...")
        
        # Configura Gemini
        api_key = os.getenv("GOOGLE_API_KEY")
        if api_key:
            genai.configure(api_key=api_key)
            self.model = genai.GenerativeModel("gemini-3.1-flash-lite")
            print("✅ Gemini configurado")
        else:
            self.model = None
            print("⚠️ GOOGLE_API_KEY não encontrada")
    
    def __call__(self, question: str) -> str:
        """Responde à pergunta com Chain of Thought."""
        
        # Se não tiver modelo, retorna N/A
        if not self.model:
            return "N/A"
        
        # Prompt melhorado com instruções claras
        prompt = f"""
        You are an AI assistant solving GAIA benchmark questions.
        
        Question: {question}
        
        Instructions:
        1. THINK STEP BY STEP about the problem.
        2. Show your reasoning briefly.
        3. END with the final answer in the format: FINAL ANSWER: [answer]
        
        Important rules:
        - If it's a number, just output the number
        - If it's a string, output just the string
        - If it's a date, output in YYYY-MM-DD format
        - No markdown, no extra text after FINAL ANSWER
        
        Let me solve this step by step:
        """
        
        try:
            response = self.model.generate_content(prompt)
            text = response.text.strip()
            
            # Extrai a resposta final
            if "FINAL ANSWER:" in text:
                answer = text.split("FINAL ANSWER:")[-1].strip()
            else:
                # Fallback: pega a última linha
                lines = [l.strip() for l in text.split('\n') if l.strip()]
                answer = lines[-1] if lines else text
            
            # Remove markdown e aspas extras
            answer = answer.strip('"').strip("'").strip()
            answer = answer.replace("```", "").strip()
            
            print(f"✅ Resposta: {answer}")
            return answer if answer else "N/A"
            
        except Exception as e:
            print(f"❌ Erro no modelo: {e}")
            return "N/A"


class SearchEnhancedAgent:
    """Agente com ferramenta de busca (se disponível)."""
    
    def __init__(self):
        print("🚀 Inicializando agente com busca...")
        
        # Configura Gemini
        api_key = os.getenv("GOOGLE_API_KEY")
        if api_key:
            genai.configure(api_key=api_key)
            self.model = genai.GenerativeModel("gemini-3.1-flash-lite")
            print("✅ Gemini configurado")
        else:
            self.model = None
            print("⚠️ GOOGLE_API_KEY não encontrada")
        
        # Tenta importar ferramentas de busca
        self.search_tool = None
        try:
            from duckduckgo_search import DDGS
            self.search_tool = DDGS()
            print("✅ DuckDuckGo configurado")
        except ImportError:
            print("⚠️ DuckDuckGo não disponível")
    
    def search(self, query: str) -> str:
        """Faz busca na web."""
        if not self.search_tool:
            return ""
        try:
            results = self.search_tool.text(query, max_results=3)
            return "\n".join([f"- {r['body']}" for r in results])
        except Exception as e:
            print(f"⚠️ Erro na busca: {e}")
            return ""
    
    def __call__(self, question: str) -> str:
        """Responde à pergunta com busca se necessário."""
        
        if not self.model:
            return "N/A"
        
        # Verifica se precisa de busca
        search_terms = ["who", "what", "when", "where", "which", "how"]
        needs_search = any(term in question.lower() for term in search_terms)
        
        search_results = ""
        if needs_search and self.search_tool:
            print("🔍 Buscando informações...")
            search_results = self.search(question)
        
        prompt = f"""
        You are an AI assistant solving GAIA benchmark questions.
        
        Question: {question}
        
        {f"Search results:\n{search_results}\n" if search_results else ""}
        
        Instructions:
        1. Use the search results if available.
        2. Think step by step.
        3. END with FINAL ANSWER: [answer]
        
        Rules:
        - Numbers: just the number
        - Strings: just the text
        - Dates: YYYY-MM-DD
        - No markdown after FINAL ANSWER
        
        Let me solve this:
        """
        
        try:
            response = self.model.generate_content(prompt)
            text = response.text.strip()
            
            if "FINAL ANSWER:" in text:
                answer = text.split("FINAL ANSWER:")[-1].strip()
            else:
                lines = [l.strip() for l in text.split('\n') if l.strip()]
                answer = lines[-1] if lines else text
            
            answer = answer.strip('"').strip("'").strip()
            answer = answer.replace("```", "").strip()
            
            print(f"✅ Resposta: {answer}")
            return answer if answer else "N/A"
            
        except Exception as e:
            print(f"❌ Erro: {e}")
            return "N/A"


# ===== FUNÇÃO PRINCIPAL =====
def run_and_submit(username: str, use_search: bool = True):
    """Executa o agente e submete as respostas."""
    
    if not username or not username.strip():
        return "❌ Digite seu username do Hugging Face.", None
    
    username = username.strip()
    print(f"\n👤 Usuário: {username}")
    
    # Escolhe o agente
    if use_search:
        agent = SearchEnhancedAgent()
    else:
        agent = ImprovedGAIAgent()
    
    # Busca perguntas
    try:
        print("📥 Buscando perguntas...")
        response = requests.get(f"{API_URL}/questions", timeout=15)
        response.raise_for_status()
        questions = response.json()
        print(f"✅ {len(questions)} perguntas carregadas")
    except Exception as e:
        return f"❌ Erro ao buscar perguntas: {e}", None
    
    # Processa perguntas
    results = []
    answers = []
    correct = 0
    
    for i, item in enumerate(questions, 1):
        task_id = item.get("task_id")
        question = item.get("question")
        
        if not task_id:
            continue
        
        print(f"\n[{i}/{len(questions)}] Task {task_id[:8]}...")
        print(f"   Pergunta: {question[:100]}...")
        
        try:
            answer = agent(question)
            answers.append({"task_id": task_id, "submitted_answer": answer})
            results.append({"Task ID": task_id, "Resposta": answer})
            print(f"   ✅ Resposta: {answer}")
        except Exception as e:
            error_msg = f"ERRO: {e}"
            results.append({"Task ID": task_id, "Resposta": error_msg})
            print(f"   ❌ {error_msg}")
    
    if not answers:
        return "❌ Nenhuma resposta gerada.", pd.DataFrame(results)
    
    # Submete
    space_id = os.getenv("SPACE_ID", "seu-usuario/seu-space")
    payload = {
        "username": username,
        "agent_code": f"https://huggingface.co/spaces/{space_id}/tree/main",
        "answers": answers
    }
    
    print(f"\n📤 Submetendo {len(answers)} respostas...")
    
    try:
        response = requests.post(f"{API_URL}/submit", json=payload, timeout=120)
        response.raise_for_status()
        data = response.json()
        
        status = f"""
✅ SUBMISSÃO CONCLUÍDA!
👤 Usuário: {data.get('username', username)}
📊 Score: {data.get('score', 'N/A')}%
✅ Acertos: {data.get('correct_count', '?')}/{data.get('total_attempted', '?')}
📝 Mensagem: {data.get('message', '')}

📋 Detalhes:
- Total perguntas: {len(questions)}
- Respostas submetidas: {len(answers)}
- Agente: {'com busca' if use_search else 'sem busca'}
"""
        return status, pd.DataFrame(results)
        
    except Exception as e:
        return f"❌ Erro na submissão: {e}", pd.DataFrame(results)


# ===== INTERFACE GRADIO =====
with gr.Blocks(title="GAIA Agent - Busca+CoT") as demo:
    gr.Markdown("""
    # 🎯 GAIA Agent - Versão Melhorada
    
    **Agente com Chain of Thought e busca na web!**
    
    Instruções:
    1. Digite seu username do Hugging Face
    2. Selecione se quer usar busca na web
    3. Clique em Executar
    4. Veja seu score!
    """)
    
    with gr.Row():
        username_input = gr.Textbox(
            label="Seu username do Hugging Face",
            placeholder="ex: nayaracardoso",
            scale=2
        )
        search_checkbox = gr.Checkbox(
            label="🔍 Usar busca na web",
            value=True
        )
        run_btn = gr.Button("🚀 Executar", variant="primary", scale=1)
    
    status_output = gr.Textbox(label="Status", lines=15)
    results_table = gr.DataFrame(label="Resultados")
    
    run_btn.click(
        fn=run_and_submit,
        inputs=[username_input, search_checkbox],
        outputs=[status_output, results_table]
    )

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