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 ketiga yang masih fresh!
FRESH_GOOGLE_KEY = "AIzaSyA2ZS7fgwNoX0RMDySxgLMUIJ6j3yZ50AQ"
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 VECTOR SEARCH DENGAN COHERE)")
    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 = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004", google_api_key=FRESH_GOOGLE_KEY)
    evaluator_embeddings = LangchainEmbeddingsWrapper(gemini_embeddings)
    
    # 2. Inisialisasi Vector Retriever (Sistem B - Cohere)
    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_system_b.csv")
    
    df_out = pd.DataFrame(columns=['user_input', 'retrieved_contexts', 'response', 'reference', 'faithfulness', 'answer_relevancy', 'context_precision'])
    if os.path.exists(out_csv):
        df_out = pd.read_csv(out_csv)
    
    for i, item in enumerate(raw_data):
        if i < len(df_out) and pd.notna(df_out.loc[i, 'faithfulness']) and pd.notna(df_out.loc[i, 'answer_relevancy']):
            continue # Skip kalau sudah dievaluasi sebelumnya (Anti-Crash)
            
        query = item.get("query")
        active_hash = item.get("active_hash")
        reference = item.get("ground_truth")
        
        logger.info(f"Mengevaluasi Baris ke-{i+1} dengan Vector Search: '{query}'")
        
        # SISTEM B: Vector Search (Hanya ambil 5 terbaik)
        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.")
            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: {e}")
            response_text = "Gagal memproses jawaban."
            
        sample = {
            "user_input": query,
            "response": response_text,
            "retrieved_contexts": contexts,
            "reference": reference if reference else "N/A"
        }
        
        dataset = EvaluationDataset.from_list([sample])
        
        try:
            result = evaluate(dataset=dataset, metrics=metrics, llm=evaluator_llm, embeddings=evaluator_embeddings, run_config=config)
            res_df = result.to_pandas()
            
            df_out.loc[i, 'user_input'] = query
            df_out.loc[i, 'retrieved_contexts'] = str(contexts)
            df_out.loc[i, 'response'] = response_text
            df_out.loc[i, 'reference'] = reference
            df_out.loc[i, 'faithfulness'] = res_df.loc[0, 'faithfulness']
            df_out.loc[i, 'answer_relevancy'] = res_df.loc[0, 'answer_relevancy']
            df_out.loc[i, 'context_precision'] = res_df.loc[0, 'context_precision']
        except Exception as e:
            logger.error(f"Error Ragas di baris {i+1}: {e}. Skipping...")
            df_out.loc[i, 'user_input'] = query
            df_out.loc[i, 'retrieved_contexts'] = str(contexts)
            df_out.loc[i, 'response'] = response_text
            df_out.loc[i, 'reference'] = reference
            df_out.loc[i, 'faithfulness'] = 0.0
            df_out.loc[i, 'answer_relevancy'] = 0.0
            df_out.loc[i, 'context_precision'] = 0.0
            
        # Simpan SEGERA setiap baris! (Anti hilang kalau crash)
        df_out.to_csv(out_csv, index=False)
        await asyncio.sleep(2) # Hindari rate limit
            
    # Round semuanya di akhir
    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__":
    import sys
    if sys.platform == 'win32':
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
    asyncio.run(main())