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

from ragas import evaluate, EvaluationDataset, RunConfig
from ragas.metrics import AnswerRelevancy, ContextPrecision, Faithfulness
from ragas.llms import LangchainLLMWrapper
from ragas.embeddings import LangchainEmbeddingsWrapper
from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings

import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from app.services.chatbot.vector_retriever import VectorRetriever

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

# Gunakan kunci Google yang fresh untuk LLM & Ragas Embeddings
FRESH_GOOGLE_KEY = "AIzaSyAAa4oW7ly9XjBjCkzJjMN3CaXHgDk9jro"
os.environ["GOOGLE_API_KEY"] = FRESH_GOOGLE_KEY

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 SISTEM B (HYBRID: COHERE RETRIEVAL + GEMINI LLM)")
    print("="*50)
    
    # 1. Inisialisasi LLM & Embeddings untuk Ragas (Google)
    gemini_llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0.0, google_api_key=FRESH_GOOGLE_KEY)
    evaluator_llm = LangchainLLMWrapper(gemini_llm)
    
    # Gemini Embeddings hanya untuk metrik AnswerRelevancy lokal Ragas
    gemini_embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001", google_api_key=FRESH_GOOGLE_KEY)
    evaluator_embeddings = LangchainEmbeddingsWrapper(gemini_embeddings)
    
    # 2. Inisialisasi Vector Retriever (Sistem B - Cohere Database)
    vector_retriever = VectorRetriever(table_name="document_chunks")
    
    metrics = [
        Faithfulness(llm=evaluator_llm),
        AnswerRelevancy(llm=evaluator_llm, embeddings=evaluator_embeddings),
        ContextPrecision(llm=evaluator_llm)
    ]
    config = RunConfig(timeout=180, max_workers=2)

    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)}")
    out_csv = os.path.join(os.path.dirname(__file__), "..", "results", "evaluation_results_hybrid.csv")
    
    df_out = pd.DataFrame(columns=['user_input', 'retrieved_contexts', 'response', 'reference', 'faithfulness', 'answer_relevancy', 'context_precision'])
    
    # Baca file jika ada untuk mode Incremental (Resume on Crash)
    if os.path.exists(out_csv):
        try:
            df_out = pd.read_csv(out_csv)
        except Exception as e:
            pass
            
    # Hapus file log lama
    if os.path.exists(out_csv) and len(df_out) == 0:
        os.remove(out_csv)
    
    # Konfigurasi Batching Ultra-Aman: 3 prompt per batch
    batch_size = 3

    for i in range(0, len(raw_data), batch_size):
        batch_num = (i // batch_size) + 1
        batch_data = raw_data[i:i+batch_size]
        
        # Cek apakah batch ini sudah dievaluasi sepenuhnya
        if i + batch_size - 1 < len(df_out):
            if pd.notna(df_out.loc[i + batch_size - 1, 'faithfulness']):
                logger.info(f"[BATCH {batch_num}] Sudah dievaluasi sebelumnya, skipping...")
                continue
        
        print("\n" + "="*50)
        logger.info(f"[BATCH {batch_num}] Memproses Soal {i+1} sampai {i+len(batch_data)}")
        
        samples = []
        
        for item in batch_data:
            query = item.get("query")
            active_hash = item.get("active_hash")
            reference = item.get("ground_truth")
            
            # SISTEM B: Vector Search Supabase menggunakan API Cohere 1024D (REST)
            chunks_info = vector_retriever.similarity_search_by_hash(query, active_hash, top_k=5)
            contexts = [c["content"] for c in chunks_info]
            
            if not contexts:
                logger.warning("Konteks kosong! Vector Search gagal/limit.")
                contexts = ["Tidak ada informasi di database."]
                
            context_str = "\n\n".join(contexts)
            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 dengan Gemini: {e}")
                response_text = "Gagal memproses jawaban."
                
            samples.append({
                "user_input": query,
                "response": response_text,
                "retrieved_contexts": contexts,
                "reference": reference if reference else "N/A"
            })
            
        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
            )
            res_df = batch_result.to_pandas()
            
            for j in range(len(res_df)):
                idx = i + j
                df_out.loc[idx, 'user_input'] = samples[j]['user_input']
                df_out.loc[idx, 'retrieved_contexts'] = str(samples[j]['retrieved_contexts'])
                df_out.loc[idx, 'response'] = samples[j]['response']
                df_out.loc[idx, 'reference'] = samples[j]['reference']
                df_out.loc[idx, 'faithfulness'] = res_df.loc[j, 'faithfulness']
                df_out.loc[idx, 'answer_relevancy'] = res_df.loc[j, 'answer_relevancy']
                df_out.loc[idx, 'context_precision'] = res_df.loc[j, 'context_precision']
                
        except Exception as e:
            logger.error(f"Error Ragas di Batch {batch_num}: {e}")
            for j in range(len(samples)):
                idx = i + j
                df_out.loc[idx, 'user_input'] = samples[j]['user_input']
                df_out.loc[idx, 'retrieved_contexts'] = str(samples[j]['retrieved_contexts'])
                df_out.loc[idx, 'response'] = samples[j]['response']
                df_out.loc[idx, 'reference'] = samples[j]['reference']
                df_out.loc[idx, 'faithfulness'] = 0.0
                df_out.loc[idx, 'answer_relevancy'] = 0.0
                df_out.loc[idx, 'context_precision'] = 0.0
            
        # Simpan SEGERA setiap batch ke CSV
        df_out.to_csv(out_csv, index=False)
        
        # Tidur 60 detik agar limit Google (15 RPM) aman!
        if i + batch_size < len(raw_data):
            logger.info("Menunggu 60 detik untuk mereset kuota Google...")
            await asyncio.sleep(60)
            
    df_out = df_out.round(2)
    df_out.to_csv(out_csv, index=False)
    print("\n[SUCCESS] SISTEM B CSV BERHASIL DISIMPAN KE:", out_csv)

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