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_cohere import ChatCohere, CohereEmbeddings

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__)

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

cohere_api_key = os.getenv("COHERE_API_KEY")

async def main():
    print("="*50)
    print("[INFO] MEMULAI EVALUASI SISTEM B (100% COHERE ENTERPRISE)")
    print("="*50)
    
    # 1. Inisialisasi Cohere untuk LLM dan Embeddings
    cohere_llm = ChatCohere(model="command-r-plus-08-2024", temperature=0.0, cohere_api_key=cohere_api_key)
    evaluator_llm = LangchainLLMWrapper(cohere_llm)
    
    cohere_emb = CohereEmbeddings(model="embed-multilingual-v3.0", cohere_api_key=cohere_api_key)
    evaluator_embeddings = LangchainEmbeddingsWrapper(cohere_emb)
    
    # 2. Inisialisasi Vector Retriever
    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)

    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'])
    
    # Batch size 3 untuk menjaga limit Cohere (20 TPM)
    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]
        
        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 Cohere
            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:
                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 = cohere_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": 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
            )
            df_batch = batch_result.to_pandas()
            
            for j in range(len(df_batch)):
                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'] = df_batch.loc[j, 'faithfulness']
                df_out.loc[idx, 'answer_relevancy'] = df_batch.loc[j, 'answer_relevancy']
                df_out.loc[idx, 'context_precision'] = df_batch.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
        df_out.to_csv(out_csv, index=False)
        
        # Tidur 60 detik agar limit Cohere 20 RPM aman!
        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__":
    import sys
    if sys.platform == 'win32':
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
    asyncio.run(main())