import os
import json
import time
import pandas as pd
from dotenv import load_dotenv

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

from langchain_cohere import ChatCohere
from langchain_groq import ChatGroq
from langchain_community.embeddings import HuggingFaceEmbeddings

from ragas import evaluate, EvaluationDataset
from ragas.metrics import AnswerRelevancy, ContextPrecision, Faithfulness
from ragas.llms import LangchainLLMWrapper
from ragas.embeddings import LangchainEmbeddingsWrapper

load_dotenv(override=True)

# Pastikan API keys tersedia
if not os.getenv("COHERE_API_KEY"):
    print("[ERROR] COHERE_API_KEY tidak ditemukan!")
    sys.exit(1)
if not os.getenv("GROQ_API_KEY"):
    print("[ERROR] GROQ_API_KEY tidak ditemukan!")
    sys.exit(1)

# Inisialisasi Retriever
v = VectorRetriever(table_name="document_chunks")

class SafeGroq(ChatGroq):
    def invoke(self, input, config=None, **kwargs):
        kwargs.pop('n', None)
        return super().invoke(input, config, **kwargs)
    
    async def ainvoke(self, input, config=None, **kwargs):
        kwargs.pop('n', None)
        return await super().ainvoke(input, config, **kwargs)

    def generate(self, messages, stop=None, callbacks=None, tags=None, metadata=None, run_name=None, run_id=None, **kwargs):
        kwargs.pop('n', None)
        return super().generate(messages, stop=stop, callbacks=callbacks, tags=tags, metadata=metadata, run_name=run_name, run_id=run_id, **kwargs)

    async def agenerate(self, messages, stop=None, callbacks=None, tags=None, metadata=None, run_name=None, run_id=None, **kwargs):
        kwargs.pop('n', None)
        return await super().agenerate(messages, stop=stop, callbacks=callbacks, tags=tags, metadata=metadata, run_name=run_name, run_id=run_id, **kwargs)

# Inisialisasi Model Penjawab (Cohere)
ans_llm = ChatCohere(model="command-r-plus-08-2024", temperature=0.0)

# Inisialisasi Model Juri (Groq)
groq_llm = SafeGroq(model="llama-3.1-8b-instant", temperature=0.0)
eval_llm = LangchainLLMWrapper(groq_llm)

# Inisialisasi Embeddings Juri Lokal (Sangat cepat dan gratis)
hf_emb = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
eval_emb = LangchainEmbeddingsWrapper(hf_emb)

metrics = [Faithfulness(llm=eval_llm), AnswerRelevancy(llm=eval_llm, embeddings=eval_emb), ContextPrecision(llm=eval_llm)]

raw_data = json.load(open('data/eval_dataset.json', 'r'))
out_csv = 'results/evaluation_results_cohere.csv'

# Resume dari CSV lama jika ada
df = pd.DataFrame(columns=['user_input', 'retrieved_contexts', 'response', 'reference', 'faithfulness', 'answer_relevancy', 'context_precision'])
if os.path.exists(out_csv):
    try:
        df = pd.read_csv(out_csv)
    except Exception:
        pass

for i, item in enumerate(raw_data):
    if i < len(df) and pd.notna(df.loc[i, 'faithfulness']):
        print(f"Skipping Soal {i+1} (Sudah ada di CSV)")
        continue
        
    print(f"Memproses Soal {i+1}/{len(raw_data)}: {item['query']}")
    
    # 1. Vector Search
    try:
        chunks = v.similarity_search_by_hash(item['query'], item['active_hash'], top_k=5)
        contexts = [c['content'] for c in chunks]
        if not contexts:
            contexts = ["Tidak ada informasi di database."]
    except Exception as e:
        print(f"  [ERROR] Vector Search: {e}")
        contexts = ["Tidak ada informasi di database."]
        
    # 2. Generate Answer with Cohere
    prompt = f"Gunakan HANYA konteks berikut untuk menjawab pertanyaan.\nKonteks:\n{chr(10).join(contexts)}\n\nPertanyaan: {item['query']}"
    try:
        ans = ans_llm.invoke(prompt).content
    except Exception as e:
        ans = "Gagal memproses."
        print(f"  [ERROR] Cohere Generate: {e}")
        # Jika Cohere limit, tunggu sebentar
        time.sleep(10)
        
    sample = [{
        "user_input": item['query'],
        "retrieved_contexts": contexts,
        "response": ans,
        "reference": item['ground_truth'] or "N/A"
    }]
    
    # 3. Evaluate with Groq
    try:
        res = evaluate(dataset=EvaluationDataset.from_list(sample), metrics=metrics, llm=eval_llm, embeddings=eval_emb, show_progress=False)
        f_score = res['faithfulness']
        ar_score = res['answer_relevancy']
        cp_score = res['context_precision']
    except Exception as e:
        print(f"  [ERROR] Ragas Groq: {e}")
        f_score, ar_score, cp_score = 0.0, 0.0, 0.0
        
    df.loc[i] = [item['query'], str(contexts), ans, item['ground_truth'], f_score, ar_score, cp_score]
    df.to_csv(out_csv, index=False)
    
    # Jeda kecil agar aman dari rate limit Groq/Cohere
    time.sleep(2)

print("EVALUASI SELESAI SECARA KESELURUHAN!")
