# scripts/benchmark_retrieval.py — Controlled Comparison & Benchmark Suite (Lexical vs E5 vs Adaptive Hybrid)
import os
import sys
import json
import time
import resource
from pathlib import Path

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from knowledge.loader import KnowledgeLoader
from knowledge.chunker import MarkdownChunker
from knowledge.index import HybridIndex
from knowledge.retrieval import HybridRetriever
from knowledge.context import ContextBuilder
from routing.query_router import QueryRouter
from policies.language import LanguageEngine

BASE_DIR = Path(__file__).resolve().parent.parent
DATA_DIR = BASE_DIR / "university_data_arabic"
EVAL_FILE = BASE_DIR / "tests" / "evaluation_questions.json"
BENCHMARK_FILE = BASE_DIR / "tests" / "retrieval_benchmark.json"


def get_ram_mb() -> float:
    """Get current resident set memory usage in MB."""
    usage = resource.getrusage(resource.RUSAGE_SELF)
    return round(usage.ru_maxrss / 1024.0, 2)


def run_benchmark():
    print("🚀 Starting Controlled Retrieval Benchmark (Lexical vs E5 vs Adaptive Hybrid)...")
    ram_start = get_ram_mb()
    print(f"📊 Initial Memory Usage: {ram_start} MB")

    # 1. Load Documents & Chunks
    t_loader_start = time.time()
    loader = KnowledgeLoader(data_dir=str(DATA_DIR))
    docs = loader.load_documents()
    chunker = MarkdownChunker()
    chunks = chunker.chunk_all(docs)
    loader_time_s = round(time.time() - t_loader_start, 2)

    print(f"📚 Loaded {len(docs)} documents into {len(chunks)} chunks in {loader_time_s}s.")

    # 2. Build HybridIndex with intfloat/multilingual-e5-small
    t_index_start = time.time()
    index = HybridIndex()
    index.build_index(chunks, enable_sentence_transformers=True)
    index_build_time_s = round(time.time() - t_index_start, 2)

    ram_after_index = get_ram_mb()
    ram_delta_mb = round(ram_after_index - ram_start, 2)
    print(f"⚡ Index build completed in {index_build_time_s}s. Memory usage: {ram_after_index} MB (+{ram_delta_mb} MB).")

    # 3. Load Evaluation Questions
    with open(EVAL_FILE, "r", encoding="utf-8") as f:
        eval_questions = json.load(f)

    retriever = HybridRetriever(index=index)
    router = QueryRouter()
    context_bld = ContextBuilder()

    modes = ["lexical_only", "e5_only", "hybrid_adaptive"]
    mode_results = {mode: [] for mode in modes}

    print(f"\n🧪 Benchmarking {len(eval_questions)} evaluation questions across 3 modes...")

    for item in eval_questions:
        q = item["question"]
        cat = item["category"]
        route = router.classify_query(q)
        intent = route["intent"]
        filter_filename = route["filter_filename"]

        for mode in modes:
            actual_mode = route.get("retrieval_mode", "hybrid") if mode == "hybrid_adaptive" else mode
            
            t0 = time.time()
            retrieved_chunks, diag = retriever.retrieve(
                query=q,
                intent=intent,
                filter_filename=filter_filename,
                retrieval_mode=actual_mode
            )
            lat_ms = round((time.time() - t0) * 1000, 2)

            ctx = context_bld.build_context(retrieved_chunks)
            unique_sources = len(set(c.filename for c in retrieved_chunks))

            res = {
                "question": q,
                "category": cat,
                "intent": intent,
                "mode": mode,
                "actual_mode_used": actual_mode,
                "latency_ms": lat_ms,
                "retrieved_chunk_count": len(retrieved_chunks),
                "unique_sources_count": unique_sources,
                "context_chars": len(ctx),
                "sources": [c.filename for c in retrieved_chunks[:5]]
            }
            mode_results[mode].append(res)

    # 4. Calculate Aggregate Metrics
    summary = {
        "dataset_info": {
            "total_documents": len(docs),
            "total_chunks": len(chunks),
            "eval_questions_count": len(eval_questions),
            "index_build_time_s": index_build_time_s,
            "ram_index_delta_mb": ram_delta_mb,
            "embedding_model": "intfloat/multilingual-e5-small"
        },
        "modes_summary": {}
    }

    print("\n==================================================")
    print("BENCHMARK COMPARISON RESULTS")
    print("==================================================")
    print(f"{'Mode':<18} | {'Avg Latency (ms)':<16} | {'Avg Sources':<12} | {'Avg Context (chars)':<20}")
    print("-" * 72)

    for mode in modes:
        res_list = mode_results[mode]
        latencies = [r["latency_ms"] for r in res_list]
        sources = [r["unique_sources_count"] for r in res_list]
        contexts = [r["context_chars"] for r in res_list]

        avg_lat = round(sum(latencies) / len(latencies), 2)
        avg_src = round(sum(sources) / len(sources), 2)
        avg_ctx = round(sum(contexts) / len(contexts), 2)

        summary["modes_summary"][mode] = {
            "avg_latency_ms": avg_lat,
            "max_latency_ms": max(latencies),
            "avg_sources_retrieved": avg_src,
            "avg_context_chars": avg_ctx
        }

        print(f"{mode:<18} | {avg_lat:<16.2f} | {avg_src:<12.2f} | {avg_ctx:<20.2f}")

    # Save to json file
    with open(BENCHMARK_FILE, "w", encoding="utf-8") as f:
        json.dump({"summary": summary, "detailed_results": mode_results}, f, ensure_ascii=False, indent=2)

    print(f"\n✅ Benchmark completed successfully! Report saved to {BENCHMARK_FILE}.")


if __name__ == "__main__":
    run_benchmark()