import os
import json
import asyncio
import logging
import pandas as pd
from dotenv import load_dotenv

# Ragas & Langchain
from ragas import evaluate, EvaluationDataset, RunConfig
from ragas.metrics import Faithfulness, AnswerRelevancy, ContextPrecision
from ragas.llms import LangchainLLMWrapper
from ragas.embeddings import LangchainEmbeddingsWrapper
from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
from supabase import create_client

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Gunakan Key Google yang masih segar
FRESH_GOOGLE_KEY = "AIzaSyAAa4oW7ly9XjBjCkzJjMN3CaXHgDk9jro"
os.environ["GOOGLE_API_KEY"] = FRESH_GOOGLE_KEY

# Pastikan path .env benar
env_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '.env'))
load_dotenv(dotenv_path=env_path, override=True)

async def main():
    print("="*50)
    print("[INFO] MEMULAI EVALUASI FINAL (GEMINI LLM + GEMINI EMBEDDINGS)")
    print("="*50)

    # 1. Inisialisasi Gemini untuk Penjawab & Penilai
    gemini_llm = ChatGoogleGenerativeAI(
        model="gemini-2.5-flash",
        temperature=0.1,
        google_api_key=FRESH_GOOGLE_KEY
    )
    evaluator_llm = LangchainLLMWrapper(gemini_llm)
    
    # 2. Inisialisasi Gemini untuk Embeddings
    gemini_embeddings = GoogleGenerativeAIEmbeddings(
        model="models/embedding-001",
        google_api_key=FRESH_GOOGLE_KEY
    )
    evaluator_embeddings = LangchainEmbeddingsWrapper(gemini_embeddings)
    
    # 3. Inisialisasi Supabase secara eksplisit (Anti-Gagal)
    supa_url = os.getenv("SUPABASE_URL")
    supa_key = os.getenv("SUPABASE_KEY")
    supabase = create_client(supa_url, supa_key)
    logger.info("Supabase terhubung.")
    
    dataset_path = os.path.join(os.path.dirname(__file__), "..", "data", "eval_dataset.json")
    with open(dataset_path, "r") as f:
        raw_data = json.load(f)

    logger.info(f"Total pertanyaan: {len(raw_data)}")
    
    # Ambil konteks HANYA SEKALI dari database untuk menghemat waktu & koneksi
    target_hash = "2b6ec4a97f7e8724693c903176db23e6"
    res = supabase.table("document_chunks").select("content").eq("file_hash", target_hash).execute()
    raw_chunks = [r["content"] for r in res.data if r.get("content")]
    logger.info(f"Berhasil menarik {len(raw_chunks)} chunks dari Supabase.")
    
    # Ambil 15 chunk teratas sebagai konteks universal untuk simulasi chatbot
    base_contexts = raw_chunks[:15]
    if not base_contexts:
        logger.error("GAGAL FATAL: CHUNKS KOSONG. Hentikan skrip.")
        return
        
    context_str = "\n\n".join(base_contexts)
    
    # Konfigurasi Batching Ultra-Aman: 3 prompt per batch
    batch_size = 3
    
    metrics = [
        Faithfulness(llm=evaluator_llm),
        AnswerRelevancy(llm=evaluator_llm, embeddings=evaluator_embeddings),
        ContextPrecision(llm=evaluator_llm)
    ]
    
    config = RunConfig(timeout=180, max_workers=2)
    all_results_dfs = []

    for i in range(0, len(raw_data), batch_size):
        batch_num = (i // batch_size) + 1
        batch_data = raw_data[i:i+batch_size]
        
        print("\n" + "="*50)
        logger.info(f"[BATCH {batch_num}] Memproses Soal {i+1} sampai {i+len(batch_data)}")
        print("="*50)
        
        samples = []
        
        for item in batch_data:
            query = item.get("query")
            reference = item.get("ground_truth")
            
            logger.info(f"-> Simulasi Chatbot Gemini untuk: '{query}'")
            
            # Simulasi Chatbot
            prompt = (
                "Anda adalah asisten peneliti. Gunakan HANYA konteks berikut untuk menjawab pertanyaan.\n"
                "Jika informasi tidak ada di konteks, jawab 'Informasi tidak ditemukan di dokumen'.\n\n"
                f"Konteks:\n{context_str}\n\n"
                f"Pertanyaan: {query}"
            )
            
            try:
                response = gemini_llm.invoke(prompt)
                response_text = response.content
            except Exception as e:
                logger.error(f"Gagal generate jawaban: {e}")
                response_text = "Gagal memproses jawaban."
            
            samples.append({
                "user_input": query,
                "response": response_text,
                "retrieved_contexts": base_contexts,
                "reference": reference
            })
            
        dataset_batch = EvaluationDataset.from_list(samples)
        
        logger.info(f"[BATCH {batch_num}] Menjalankan Penilaian Ragas...")
        try:
            batch_result = evaluate(
                dataset=dataset_batch,
                metrics=metrics,
                llm=evaluator_llm,
                embeddings=evaluator_embeddings,
                run_config=config
            )
            df_batch = batch_result.to_pandas()
            all_results_dfs.append(df_batch)
        except Exception as e:
            logger.error(f"[ERROR] Penilaian Ragas gagal di Batch {batch_num}: {e}")
            
        if i + batch_size < len(raw_data):
            logger.info(f"[JEDA BATCH] Menunggu 60 detik agar Kuota 15 RPM Google ter-reset penuh...")
            await asyncio.sleep(60)
            
    if all_results_dfs:
        final_df = pd.concat(all_results_dfs, ignore_index=True)
        final_df = final_df.round(2)
        
        print("\n" + "="*50)
        print("[RESULT] HASIL EVALUASI FINAL (ULTIMATE)")
        print("="*50)
        print(final_df[['user_input', 'faithfulness', 'answer_relevancy', 'context_precision']].to_markdown())
        
        results_dir = os.path.join(os.path.dirname(__file__), "..", "results")
        os.makedirs(results_dir, exist_ok=True)
        out_path = os.path.join(results_dir, "evaluation_results_ultimate.csv")
        final_df.to_csv(out_path, index=False)
        logger.info(f"[SUCCESS] Laporan disimpan ke: {out_path}")

if __name__ == "__main__":
    import sys
    if sys.platform == 'win32':
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
    asyncio.run(main())
