from llama_index.core.evaluation import CorrectnessEvaluator
from ketju.rag.llama_index.variants.simple import SimplePdfRag
from ketju.rag.llama_index.variants.full_doc_in_ctx import FullDocInCtxRag
from ketju.config import DOCS_DIR
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import Settings
from typing import Tuple, List, Dict
from pathlib import Path
from llama_index.core.llama_dataset import (
    LabelledRagDataset,
    CreatedBy,
    CreatedByType,
    LabelledRagDataExample,
)
#file = DOCS_DIR / "wärtsilä/energy-docs/technology-products/product-guide/W32 PG.pdf"
import pickle
from devtools import debug

Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
Settings.llm = OpenAI(model="gpt-4o-mini")

def get_large_dataset() -> Tuple[List[Path], List[Dict[str, str]]]:
    base_dir = DOCS_DIR / "interdimensional-systems"
    eval_file = base_dir / "eval.json"

    # Load evaluation JSON
    import json
    with open(eval_file, "r") as f:
        eval_json = json.load(f)

    # Collect all PDF files in the folder
    files = list(base_dir.glob("*.pdf"))

    # Convert JSON questions to expected dataset format
    dataset = [
        {
            "query": q["query"],
            "correct_answer": q["answer"],
        }
        for q in eval_json["questions"]
    ]

    return files, dataset




def evaluate():
    files, dataset = get_large_dataset()

    evaluator = CorrectnessEvaluator(
        llm=OpenAI(model="gpt-4o-mini", temperature=0),
        score_threshold=4.0
    )

    instances = {
        "simple_rag": SimplePdfRag(name="simple_rag"),
        "full_doc_in_ctx_rag": FullDocInCtxRag(name="full_doc_in_ctx_rag"),
    }

    evaluation_data = {}
    for name, rag in instances.items():
        print(f"\nEvaluating {name}")
        print("=" * 40)
        evaluation_data[name] = []

        rag.ingest(files=files)

        for sample in dataset:
            response = rag.query(sample["query"])
            response_text = str(response).strip()

            evaluation = evaluator.evaluate(
                query=sample["query"],
                response=response_text,
                reference=sample["correct_answer"]
            )
            result_data = {
                "rag_engine": name,
                "model": str(Settings.llm.model),
                "query": sample["query"],
                "response": response_text,
                "correct_answer": sample["correct_answer"],
                "score": evaluation.score,
                "passing": evaluation.passing,
                "feedback": evaluation.feedback,
            }
            evaluation_data[name].append(result_data)
            debug(result_data)
            


    with open("evaluation_data.pkl", "wb") as f:
        pickle.dump(evaluation_data, f)
    return evaluation_data

def load_evaluation_data():
    with open("evaluation_data.pkl", "rb") as f:
        return pickle.load(f)

print("\nENGINE COMPARISON SUMMARY")
print("=" * 40)

if Path ("evaluation_data.pkl").exists():
    evaluation_data = load_evaluation_data()
else:
    evaluation_data = evaluate()

for name, results in evaluation_data.items():
    passed = sum(result["passing"] for result in results)
    total = len(results)
    avg_score = sum(result["score"] for result in results) / total if total > 0 else 0.0

    print(
        f"{name}: "
        f"Pass rate = {passed}/{total}, "
        f"Avg score = {avg_score:.2f}"
    )
