# ==============================================================================
#  AetherMap RAG Evaluation Script
#  Avalia qualidade do sistema RAG usando RAGAS (sem ground-truth)
# ==============================================================================
"""
Uso:
    python evaluate_rag.py --api-url https://madras1-aethermap.hf.space --dataset data.csv --text-column texto

Métricas avaliadas (sem precisar de ground-truth):
    - Faithfulness: A resposta é fiel ao contexto recuperado?
    - Answer Relevancy: A resposta é relevante à pergunta?
"""

import os
import json
import argparse
import requests
from typing import List, Dict, Any
from dataclasses import dataclass
from datetime import datetime

# Avaliação
from ragas import evaluate
from ragas.metrics import faithfulness, answer_relevancy
from datasets import Dataset

# LLM para avaliação (Groq via LangChain)
from langchain_groq import ChatGroq


@dataclass
class RAGSample:
    """Representa uma amostra para avaliação RAG."""
    question: str
    answer: str
    contexts: List[str]


class AetherMapEvaluator:
    """Avaliador de RAG para o AetherMap."""
    
    def __init__(self, api_url: str, groq_api_key: str = None):
        self.api_url = api_url.rstrip("/")
        self.groq_api_key = groq_api_key or os.environ.get("GROQ_API_KEY")
        self.job_id = None
        
        if not self.groq_api_key:
            raise ValueError("GROQ_API_KEY não encontrada. Defina via argumento ou variável de ambiente.")
        
        # LLM avaliador (diferente do modelo de geração)
        self.eval_llm = ChatGroq(
            model="llama-3.3-70b-versatile",  # Modelo forte para avaliação
            api_key=self.groq_api_key,
            temperature=0
        )
    
    def upload_dataset(self, file_path: str, text_column: str = None, n_samples: int = 1000) -> bool:
        """Faz upload de um dataset para o AetherMap."""
        print(f"📤 Uploading dataset: {file_path}")
        
        with open(file_path, "rb") as f:
            files = {"file": (os.path.basename(file_path), f)}
            data = {"n_samples": n_samples}
            
            if text_column and file_path.lower().endswith(".csv"):
                data["text_column"] = text_column
            
            response = requests.post(f"{self.api_url}/process/", files=files, data=data)
        
        if response.status_code != 200:
            print(f"❌ Erro no upload: {response.text}")
            return False
        
        result = response.json()
        self.job_id = result["job_id"]
        print(f"✅ Dataset carregado! Job ID: {self.job_id}")
        print(f"   📊 {result['metadata']['num_documents_processed']} documentos")
        print(f"   🎯 {result['metadata']['num_clusters_found']} clusters")
        return True
    
    def search(self, query: str) -> RAGSample:
        """Executa uma busca e retorna os componentes RAG."""
        if not self.job_id:
            raise ValueError("Nenhum dataset carregado. Use upload_dataset primeiro.")
        
        response = requests.post(
            f"{self.api_url}/search/",
            data={"query": query, "job_id": self.job_id}
        )
        
        if response.status_code != 200:
            raise Exception(f"Erro na busca: {response.text}")
        
        result = response.json()
        
        # Extrair contextos (precisamos buscar os textos completos)
        contexts = []
        for res in result.get("results", []):
            # O resultado inclui o índice, precisamos do texto
            # Por enquanto, vamos usar o summary que inclui as citações
            pass
        
        # Como não temos acesso direto aos textos via API,
        # vamos extrair do summary as partes entre [ID: x]
        # Ou podemos modificar a API para retornar os contextos
        
        return RAGSample(
            question=query,
            answer=result.get("summary", ""),
            contexts=contexts if contexts else [result.get("summary", "")]  # Fallback
        )
    
    def search_with_contexts(self, query: str) -> RAGSample:
        """
        Executa busca e extrai contextos completos.
        Requer endpoint modificado ou acesso ao cache.
        """
        if not self.job_id:
            raise ValueError("Nenhum dataset carregado.")
        
        # Primeira chamada: buscar documentos relevantes
        response = requests.post(
            f"{self.api_url}/search/",
            data={"query": query, "job_id": self.job_id}
        )
        
        if response.status_code != 200:
            raise Exception(f"Erro: {response.text}")
        
        result = response.json()
        answer = result.get("summary", "")
        
        # Extrair contextos dos resultados
        # Nota: Idealmente o endpoint deveria retornar os textos completos
        # Por agora, usamos os índices para referência
        contexts = []
        for res in result.get("results", []):
            # Placeholder - em produção, você buscaria o texto real
            contexts.append(f"Documento #{res['index']} (score: {res['score']:.3f})")
        
        # Se não conseguiu contextos, usa o summary como fallback
        if not contexts:
            contexts = [answer]
        
        return RAGSample(
            question=query,
            answer=answer,
            contexts=contexts
        )
    
    def evaluate_queries(self, queries: List[str], verbose: bool = True) -> Dict[str, Any]:
        """
        Avalia uma lista de queries usando RAGAS.
        
        Métricas:
            - faithfulness: Resposta é fiel ao contexto?
            - answer_relevancy: Resposta é relevante à pergunta?
        """
        print(f"\n🔍 Avaliando {len(queries)} queries...")
        
        samples = []
        for i, query in enumerate(queries):
            if verbose:
                print(f"   [{i+1}/{len(queries)}] {query[:50]}...")
            
            try:
                sample = self.search_with_contexts(query)
                samples.append(sample)
            except Exception as e:
                print(f"   ⚠️ Erro na query {i+1}: {e}")
        
        if not samples:
            return {"error": "Nenhuma amostra válida coletada."}
        
        # Preparar dataset para RAGAS
        eval_data = {
            "question": [s.question for s in samples],
            "answer": [s.answer for s in samples],
            "contexts": [s.contexts for s in samples],
        }
        
        dataset = Dataset.from_dict(eval_data)
        
        print("\n📊 Executando avaliação RAGAS...")
        
        # Avaliar com RAGAS
        try:
            result = evaluate(
                dataset,
                metrics=[faithfulness, answer_relevancy],
                llm=self.eval_llm,
            )
            
            return {
                "metrics": {
                    "faithfulness": float(result["faithfulness"]),
                    "answer_relevancy": float(result["answer_relevancy"]),
                },
                "num_samples": len(samples),
                "timestamp": datetime.now().isoformat(),
            }
        except Exception as e:
            return {"error": str(e)}
    
    def generate_test_queries(self, n: int = 10) -> List[str]:
        """
        Gera queries de teste baseadas no dataset carregado.
        Usa o LLM para criar perguntas relevantes.
        """
        # Queries genéricas de exemplo
        # Em produção, você geraria baseado no conteúdo real
        return [
            "Qual é o tema principal dos documentos?",
            "Quais são as entidades mais mencionadas?",
            "Existe alguma tendência nos dados?",
            "Quais são os pontos mais controversos?",
            "Qual a conclusão principal?",
            "Existem contradições entre os documentos?",
            "Qual documento é mais relevante para o tema central?",
            "Quais termos técnicos são mais usados?",
            "Há alguma relação temporal entre os eventos?",
            "Qual é o sentimento geral dos textos?",
        ][:n]


def main():
    parser = argparse.ArgumentParser(description="Avaliador RAG para AetherMap")
    parser.add_argument("--api-url", default="https://madras1-aethermap.hf.space",
                        help="URL da API do AetherMap")
    parser.add_argument("--dataset", help="Arquivo de dataset para upload (CSV ou TXT)")
    parser.add_argument("--text-column", help="Coluna de texto para CSV")
    parser.add_argument("--job-id", help="Job ID existente (pula upload)")
    parser.add_argument("--queries", nargs="+", help="Queries para testar")
    parser.add_argument("--queries-file", help="Arquivo com queries (uma por linha)")
    parser.add_argument("--n-samples", type=int, default=1000, help="Número de amostras do dataset")
    parser.add_argument("--output", help="Arquivo de saída para resultados (JSON)")
    
    args = parser.parse_args()
    
    # Inicializar avaliador
    evaluator = AetherMapEvaluator(api_url=args.api_url)
    
    # Upload dataset ou usar job existente
    if args.job_id:
        evaluator.job_id = args.job_id
        print(f"📎 Usando Job ID existente: {args.job_id}")
    elif args.dataset:
        if not evaluator.upload_dataset(args.dataset, args.text_column, args.n_samples):
            return
    else:
        print("❌ Forneça --dataset ou --job-id")
        return
    
    # Carregar queries
    queries = []
    if args.queries:
        queries = args.queries
    elif args.queries_file:
        with open(args.queries_file, "r", encoding="utf-8") as f:
            queries = [line.strip() for line in f if line.strip()]
    else:
        print("⚡ Gerando queries de teste...")
        queries = evaluator.generate_test_queries(10)
    
    print(f"📝 {len(queries)} queries para avaliar")
    
    # Executar avaliação
    results = evaluator.evaluate_queries(queries)
    
    # Mostrar resultados
    print("\n" + "="*50)
    print("📊 RESULTADOS DA AVALIAÇÃO RAGAS")
    print("="*50)
    
    if "error" in results:
        print(f"❌ Erro: {results['error']}")
    else:
        metrics = results["metrics"]
        print(f"\n🎯 Faithfulness:      {metrics['faithfulness']:.3f}")
        print(f"🎯 Answer Relevancy:  {metrics['answer_relevancy']:.3f}")
        print(f"\n📈 Amostras avaliadas: {results['num_samples']}")
        
        # Interpretação
        avg_score = (metrics['faithfulness'] + metrics['answer_relevancy']) / 2
        if avg_score >= 0.8:
            print("\n✅ Excelente! O RAG está performando muito bem.")
        elif avg_score >= 0.6:
            print("\n🟡 Bom, mas há espaço para melhorias.")
        else:
            print("\n🔴 Precisa de melhorias significativas.")
    
    # Salvar resultados
    if args.output:
        with open(args.output, "w", encoding="utf-8") as f:
            json.dump(results, f, indent=2, ensure_ascii=False)
        print(f"\n💾 Resultados salvos em: {args.output}")


if __name__ == "__main__":
    main()