"""
RAGAS evaluation script.
Dùng LLM-as-judge thay vì ragas library để tránh dependency conflicts.
Metrics: Faithfulness, Answer Relevancy, Context Recall
"""
print("Loading evaluate.py...")
import json
print("json ok")
import time
print("time ok")
from groq import Groq
print("groq ok")
from dotenv import load_dotenv
import os
load_dotenv()
client = Groq(api_key=os.environ["GROQ_API_KEY"])
JUDGE_MODEL = os.getenv("EVAL_MODEL", "llama-3.1-8b-instant")
print("groq client ok")
import src.rag as rag_module
rag_module.LLM_MODEL = JUDGE_MODEL
from src.rag import answer as rag_answer
print("rag ok")

CHECKPOINT_PATH = "eval_checkpoint.json"


def _save_checkpoint(results: list):
    with open(CHECKPOINT_PATH, "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)


def _load_checkpoint() -> list:
    if os.path.exists(CHECKPOINT_PATH):
        with open(CHECKPOINT_PATH, encoding="utf-8") as f:
            data = json.load(f)
        print(f"Tiếp tục từ checkpoint: {len(data)} câu đã xong.")
        return data
    return []


def judge(prompt: str) -> float:
    """Gọi LLM judge, trả về score 0.0-1.0."""
    response = client.chat.completions.create(
        model=JUDGE_MODEL,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=10,
        temperature=0.0,
    )
    text = response.choices[0].message.content.strip()
    try:
        score = float(text.split()[0])
        return max(0.0, min(1.0, score))
    except Exception:
        return 0.5


def eval_faithfulness(question: str, answer: str, contexts: list[str]) -> float:
    context_text = "\n".join(c[:600] for c in contexts[:3])  # top-3, max 600 chars mỗi chunk
    prompt = f"""Đánh giá xem câu trả lời có trung thực với ngữ cảnh không.
Chỉ trả về một số từ 0.0 đến 1.0 (1.0 = hoàn toàn trung thực, 0.0 = bịa đặt).

Ngữ cảnh: {context_text}
Câu trả lời: {answer}

Score:"""
    return judge(prompt)


def eval_answer_relevancy(question: str, answer: str) -> float:
    prompt = f"""Đánh giá xem câu trả lời có liên quan đến câu hỏi không.
Chỉ trả về một số từ 0.0 đến 1.0 (1.0 = rất liên quan, 0.0 = không liên quan).

Câu hỏi: {question}
Câu trả lời: {answer}

Score:"""
    return judge(prompt)


def eval_context_recall(question: str, ground_truth: str, contexts: list[str]) -> float:
    context_text = "\n".join(c[:600] for c in contexts[:3])
    prompt = f"""Đánh giá xem ngữ cảnh có chứa đủ thông tin để trả lời câu hỏi theo đáp án chuẩn không.
Chỉ trả về một số từ 0.0 đến 1.0 (1.0 = đủ thông tin, 0.0 = thiếu hoàn toàn).

Câu hỏi: {question}
Đáp án chuẩn: {ground_truth}
Ngữ cảnh: {context_text}

Score:"""
    return judge(prompt)


def _print_summary(results: list):
    avg_faith = sum(r["faithfulness"] for r in results) / len(results)
    avg_rel = sum(r["answer_relevancy"] for r in results) / len(results)
    avg_recall = sum(r["context_recall"] for r in results) / len(results)

    print(f"\n{'='*50}")
    print(f"KẾT QUẢ ĐÁNH GIÁ ({len(results)} câu hỏi)")
    print(f"Judge model: {JUDGE_MODEL}")
    print(f"{'='*50}")
    print(f"Faithfulness:     {avg_faith:.3f}")
    print(f"Answer Relevancy: {avg_rel:.3f}")
    print(f"Context Recall:   {avg_recall:.3f}")
    print(f"Average:          {(avg_faith + avg_rel + avg_recall) / 3:.3f}")
    return avg_faith, avg_rel, avg_recall


def evaluate(golden_set_path: str = "data/golden_set.json", sample: int = 20):
    with open(golden_set_path, encoding="utf-8") as f:
        golden = json.load(f)

    items = [q for q in golden if q.get("type") != "out_of_scope"][:sample]

    # Resume từ checkpoint nếu có
    results = _load_checkpoint()
    done_questions = {r["question"] for r in results}
    remaining = [it for it in items if it["question"] not in done_questions]

    print(f"Đánh giá {len(items)} câu hỏi ({len(remaining)} còn lại)...\n")

    for i, item in enumerate(remaining, len(results) + 1):
        question = item["question"]
        ground_truth = item["answer"]

        print(f"[{i}/{len(items)}] {question[:60]}...")
        try:
            result = rag_answer(question, top_k=5)
            ans = result["answer"]
            context_texts = result["contexts"]

            faith = eval_faithfulness(question, ans, context_texts)
            relevancy = eval_answer_relevancy(question, ans)
            recall = eval_context_recall(question, ground_truth, context_texts)

            results.append({
                "question": question,
                "ground_truth": ground_truth,
                "answer": ans,
                "faithfulness": faith,
                "answer_relevancy": relevancy,
                "context_recall": recall,
            })

            print(f"  Faithfulness: {faith:.2f} | Relevancy: {relevancy:.2f} | Recall: {recall:.2f}")
            _save_checkpoint(results)
            time.sleep(2)

        except Exception as e:
            print(f"  Lỗi: {e}")
            time.sleep(5)

    if results:
        avg_faith, avg_rel, avg_recall = _print_summary(results)

        with open("eval_results.json", "w", encoding="utf-8") as f:
            json.dump({
                "summary": {
                    "faithfulness": avg_faith,
                    "answer_relevancy": avg_rel,
                    "context_recall": avg_recall,
                    "judge_model": JUDGE_MODEL,
                },
                "details": results,
            }, f, ensure_ascii=False, indent=2)
        print("\nĐã lưu chi tiết vào eval_results.json")

        # Xóa checkpoint khi hoàn thành
        if len(results) == len(items) and os.path.exists(CHECKPOINT_PATH):
            os.remove(CHECKPOINT_PATH)
            print("Đã xóa checkpoint.")


if __name__ == "__main__":
    import sys
    golden_path = sys.argv[1] if len(sys.argv) > 1 else "data/golden_set.json"
    sample = int(sys.argv[2]) if len(sys.argv) > 2 else 20
    evaluate(golden_set_path=golden_path, sample=sample)