# ==============================================================================
#  AetherMap RAG Evaluation v3 - OpenRouter Edition
#  Usa OpenRouter ao invés de Groq para evitar rate limits
# ==============================================================================

# %% 
# Célula 1: Instalação 
# !pip install requests pandas openai -q

# %%
# Célula 2: Imports
import os
import json
import requests
import pandas as pd
import numpy as np
from typing import List, Dict, Any, Tuple
from dataclasses import dataclass, field
import tempfile
import time
import random
from openai import OpenAI

# Configurar API Key
OPENROUTER_API_KEY = "sk-or-v1-..."  # 👈 Sua key OpenRouter
API_URL = "https://madras1-aethermap.hf.space"

# %%
# Célula 3: Cliente OpenRouter
client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENROUTER_API_KEY,
)

# Modelos gratuitos do OpenRouter
QUERY_GENERATOR_MODEL = "meta-llama/llama-3.1-8b-instruct:free"
EVALUATOR_MODEL = "google/gemini-2.0-flash-exp:free"

# Rate limiting mais leve (OpenRouter é mais generoso)
DELAY_BETWEEN_CALLS = 1.5  # Segundos entre chamadas

print(f"🤖 Gerador: {QUERY_GENERATOR_MODEL}")
print(f"🧑‍⚖️ Avaliador: {EVALUATOR_MODEL}")


# %%
# Célula 4: Funções

def call_llm(messages: List[Dict], model: str, max_tokens: int = 1000) -> str:
    """Chama OpenRouter API."""
    time.sleep(DELAY_BETWEEN_CALLS)
    
    try:
        completion = client.chat.completions.create(
            model=model,
            messages=messages,
            max_tokens=max_tokens,
            temperature=0.3,
        )
        return completion.choices[0].message.content.strip()
    except Exception as e:
        print(f"   ⚠️ Erro LLM: {e}")
        time.sleep(5)  # Esperar se der erro
        return ""


def gerar_queries(documentos: List[str], n_queries: int = 5) -> List[Dict[str, str]]:
    """Gera perguntas baseadas nos documentos."""
    
    # Limitar tamanho para não estourar contexto
    docs_text = "\n---\n".join([doc[:400] for doc in documentos[:3]])
    
    prompt = f"""Gere {n_queries} perguntas que podem ser respondidas com estes documentos.
Para cada pergunta, dê também a resposta correta.

DOCUMENTOS:
{docs_text}

Retorne APENAS JSON:
[{{"pergunta": "...", "resposta": "..."}}]"""
    
    response = call_llm(
        [{"role": "user", "content": prompt}],
        QUERY_GENERATOR_MODEL,
        max_tokens=1000
    )
    
    try:
        if "[" in response:
            json_start = response.index("[")
            json_end = response.rindex("]") + 1
            return json.loads(response[json_start:json_end])[:n_queries]
    except:
        pass
    return []


def avaliar_resposta(pergunta: str, resposta_rag: str, resposta_esperada: str) -> Dict[str, float]:
    """Avalia resposta do RAG."""
    
    prompt = f"""Avalie a resposta do sistema comparando com a esperada.

PERGUNTA: {pergunta}
RESPOSTA ESPERADA: {resposta_esperada}
RESPOSTA DO SISTEMA: {resposta_rag[:800]}

Dê uma nota de 0 a 1 para:
- correcao: informação está correta?
- completude: resposta está completa?

Retorne APENAS JSON: {{"correcao": X.X, "completude": X.X}}"""
    
    response = call_llm(
        [{"role": "user", "content": prompt}],
        EVALUATOR_MODEL,
        max_tokens=200
    )
    
    try:
        if "{" in response:
            json_start = response.index("{")
            json_end = response.rindex("}") + 1
            result = json.loads(response[json_start:json_end])
            return {
                "correcao": float(result.get("correcao", 0.5)),
                "completude": float(result.get("completude", 0.5))
            }
    except:
        pass
    return {"correcao": 0.5, "completude": 0.5}


# %%
# Célula 5: Classe do avaliador

class RAGEvaluator:
    def __init__(self, api_url: str = API_URL):
        self.api_url = api_url.rstrip("/")
        print("✅ Avaliador OpenRouter inicializado!")
    
    def load_dataset(self, file_path: str, text_col: str, label_col: str):
        with open(file_path, 'r', encoding='utf-8') as f:
            sep = ';' if ';' in f.readline() else ','
        df = pd.read_csv(file_path, sep=sep)
        print(f"📂 {len(df)} linhas carregadas")
        
        splits = {}
        for label in sorted(df[label_col].unique()):
            splits[int(label)] = df[df[label_col] == label]
            print(f"   Label {label}: {len(splits[int(label)])}")
        return splits
    
    def upload(self, df_subset, text_col: str, n_samples: int = 200):
        sample = df_subset.sample(n=min(n_samples, len(df_subset)), random_state=42)
        sample_docs = sample[text_col].tolist()[:10]
        
        with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False, encoding='utf-8') as f:
            sample.to_csv(f, index=False)
            temp_path = f.name
        
        print(f"   📤 Upload {len(sample)} docs...")
        with open(temp_path, "rb") as f:
            response = requests.post(
                f"{self.api_url}/process/",
                files={"file": ("data.csv", f)},
                data={"n_samples": n_samples, "text_column": text_col, "fast_mode": "true"}
            )
        os.unlink(temp_path)
        
        if response.status_code != 200:
            print(f"   ❌ {response.text[:100]}")
            return None, []
        
        result = response.json()
        print(f"   ✅ Job: {result['job_id']}")
        return result['job_id'], sample_docs
    
    def search(self, job_id: str, query: str) -> str:
        response = requests.post(
            f"{self.api_url}/search/",
            data={"query": query, "job_id": job_id}
        )
        if response.status_code != 200:
            return "Erro"
        return response.json().get("summary", "")
    
    def evaluate_label(self, job_id: str, docs: List[str], label: int, n_queries: int = 3):
        print(f"\n📊 LABEL {label}")
        
        # Gerar queries
        queries = gerar_queries(docs, n_queries)
        if not queries:
            return {"error": "Falha queries", "label": label}
        
        print(f"   {len(queries)} queries geradas")
        
        # Testar
        scores = []
        for i, q in enumerate(queries):
            pergunta = q.get("pergunta", q.get("question", ""))
            esperada = q.get("resposta", q.get("answer", q.get("resposta_esperada", "")))
            
            print(f"   [{i+1}] {pergunta[:40]}...")
            
            # RAG responde
            resposta = self.search(job_id, pergunta)
            
            # Avaliar
            aval = avaliar_resposta(pergunta, resposta, esperada)
            scores.append((aval["correcao"] + aval["completude"]) / 2)
            
            emoji = "✓" if scores[-1] >= 0.5 else "○"
            print(f"       {emoji} {aval['correcao']:.2f} / {aval['completude']:.2f}")
        
        return {
            "label": label,
            "score": float(np.mean(scores)),
            "std": float(np.std(scores)),
            "n": len(scores)
        }
    
    def run(self, file_path: str, text_col: str = "texto", 
            label_col: str = "label_qualidade", 
            samples: int = 200, queries: int = 3):
        
        print("="*50)
        print("🌌 AETHERMAP RAG EVAL v3 (OpenRouter)")
        print("="*50)
        
        splits = self.load_dataset(file_path, text_col, label_col)
        results = {}
        
        for label, df_subset in splits.items():
            job_id, docs = self.upload(df_subset, text_col, samples)
            if not job_id:
                results[label] = {"error": "Upload falhou"}
                continue
            
            time.sleep(3)
            results[label] = self.evaluate_label(job_id, docs, label, queries)
        
        # Resultado
        print("\n" + "="*50)
        print("📊 RESULTADO FINAL")
        print("="*50)
        for label, data in sorted(results.items()):
            if "error" in data:
                print(f"Label {label}: ❌ {data['error']}")
            else:
                emoji = "✅" if data['score'] >= 0.6 else "🟡" if data['score'] >= 0.4 else "🔴"
                print(f"Label {label}: {data['score']:.3f} ±{data['std']:.2f} {emoji}")
        
        return results


# %%
# Executar
evaluator = RAGEvaluator()

# %%
resultados = evaluator.run(
    file_path="dataset_unificado_72k.csv",
    text_col="texto",
    label_col="label_qualidade",
    samples=200,   # Menos samples
    queries=3      # Menos queries
)