#!/usr/bin/env python3
"""Bewertet die Retrieval-Qualität gegen einen Satz Referenzfragen.

Das Skript stellt jede Frage aus eval/goldset.csv an den Vektorspeicher und prüft,
ob die erwartete Fundstelle unter den ersten k Treffern auftaucht. Ergebnis sind
Recall@k und MRR, ausgegeben auf dem Bildschirm und in eval/baseline.txt.

Optional (--judge MODELL) beurteilt zusätzlich ein Sprachmodell über die OpenAI-API,
ob die Antwort in den ersten JUDGE_K Treffern steht — gleich in welchem Buch und auf
welcher Seite. Das misst, was für den Chat zählt, und wertet gleichwertige Stellen
aus anderen Quellen nicht als Fehler. Braucht OPENAI_API_KEY; die Urteile werden in
eval/judge_cache.json zwischengespeichert, damit Wiederholungen nichts kosten.

Aus dem Projektordner starten:  python eval_retrieval.py

Fragen, Index und Bericht lassen sich auf der Kommandozeile überschreiben, z. B.
    python eval_retrieval.py --goldset eval/goldset_kafka.csv --index ./vectorstore_docling \
        --report eval/kafka_docling.txt
"""
import argparse
import csv
import sys
import unicodedata
from datetime import datetime
from pathlib import Path

from local_rag import (
    CHUNK_OVERLAP,
    CHUNK_SIZE,
    QUERY_ADDITIONAL_INSTRUCTIONS,
    LocalRAG,
)

# --- Einstellungen -----------------------------------------------------------

# Datei mit den Referenzfragen (Spalten siehe eval/goldset.csv).
GOLDSET_PATH = Path("eval/goldset.csv")

# Wohin der Bericht geschrieben wird. Diese Datei gehört ins Repository.
REPORT_PATH = Path("eval/baseline.txt")

# Muss mit server.py übereinstimmen, sonst wird ein anderer Index gemessen.
VECTORSTORE_PATH = "./vectorstore_pypdf_2000"

# Für diese k-Werte wird Recall@k berechnet.
K_VALUES = [5, 10, 20]

# Den Zusatztext aus local_rag.QUERY_ADDITIONAL_INSTRUCTIONS an jede Frage anhängen?
# Einmal mit True und einmal mit False laufen lassen zeigt, ob er tatsächlich hilft.
USE_QUERY_INSTRUCTIONS = False

# PyPDFLoader zählt Seiten ab 0, Menschen ab 1. Bei 1 bedeutet Seite 42 in der
# CSV also die Seite, die im PDF-Betrachter als 42 angezeigt wird.
PAGE_OFFSET = 1

# Ein Chunk kann über einen Seitenumbruch reichen: so viele Seiten Abweichung
# gelten noch als Treffer. 0 verlangt die exakte Seite.
PAGE_TOLERANCE = 1

# Wie viele Treffer das Sprachmodell bei --judge zu sehen bekommt (wie RETRIEVAL_K im Chat).
JUDGE_K = 5
JUDGE_CACHE = Path("eval/judge_cache.json")

JUDGE_PROMPT = """Du prüfst eine Literatursuche zu Franz Kafka. Unten stehen eine Frage und \
{n} nummerierte Textauszüge aus Primär- und Sekundärliteratur (teils mit OCR-Fehlern).

Entscheide für jeden Auszug, ob er die Frage beantwortet: ganz oder in ihrem wesentlichen \
Kern, so dass man allein mit diesem Auszug eine zutreffende Antwort geben könnte. Ein Auszug, \
der nur das Thema berührt, zählt nicht. Nutze kein eigenes Wissen über Kafka, nur die Auszüge.

Antworte ausschließlich mit JSON in dieser Form:
{{"auszuege": [Nummern der Auszüge, die die Frage beantworten], "begruendung": "ein Satz"}}

Frage: {frage}

{auszuege}"""

RETRIEVAL_SETUP = ""   # wird in main() gesetzt und im Bericht vermerkt
INDEX_INFO = {}        # Einstellungen, mit denen der Index gebaut wurde (index_info.json)
JUDGE_MODEL = None     # wird in main() aus --judge gesetzt

# -----------------------------------------------------------------------------


def read_goldset(path):
    """Liest die Referenzfragen. Erwartet die Spalten frage, korpus,
    erwartete_quelle, erwartete_seite, anmerkung (die letzten beiden optional)."""
    if not path.exists():
        sys.exit(
            f"Keine Referenzfragen gefunden: {path}\n"
            "Lege die Datei an (Vorlage liegt im Repository) und trage echte Fragen ein."
        )

    rows = []
    with path.open(encoding="utf-8-sig", newline="") as handle:
        # Deutsches Excel speichert CSV oft mit Semikolon statt Komma.
        first_line = handle.readline()
        delimiter = ";" if first_line.count(";") > first_line.count(",") else ","
        handle.seek(0)

        for line_no, row in enumerate(csv.DictReader(handle, delimiter=delimiter), start=2):
            question = (row.get("frage") or "").strip()
            source = (row.get("erwartete_quelle") or "").strip()
            if not question or not source:
                print(f"  Zeile {line_no} übersprungen: frage oder erwartete_quelle fehlt")
                continue

            page_raw = (row.get("erwartete_seite") or "").strip()
            page = None
            if page_raw:
                try:
                    page = int(page_raw)
                except ValueError:
                    print(f"  Zeile {line_no}: '{page_raw}' ist keine Seitenzahl, wird ignoriert")

            rows.append({
                "question": question,
                "corpus": (row.get("korpus") or "ohne").strip() or "ohne",
                "source": source,
                "page": page,
            })

    if not rows:
        sys.exit(f"{path} enthält keine verwertbaren Zeilen.")
    return rows


def same_file(found, expected):
    """Gleiche Datei? Groß-/Kleinschreibung, Unicode-Form (macOS speichert Umlaute zerlegt)
    und Unterstriche, Punkte oder Leerzeichen am Ende des Namens spielen keine Rolle."""
    def norm(name):
        stem = unicodedata.normalize("NFC", Path(name).stem.lower())
        return stem.rstrip("_. ")
    return norm(found) == norm(expected)


def is_hit(doc, expected_source, expected_page):
    """Passt dieser Treffer zur erwarteten Fundstelle? Dateiname ohne Pfad,
    Groß-/Kleinschreibung egal; Seite nur, wenn eine angegeben wurde."""
    if not same_file(doc.metadata.get("source", ""), expected_source):
        return False

    if expected_page is None:
        return True

    found_page = doc.metadata.get("page")
    if found_page is None:
        return False
    return abs(int(found_page) - (expected_page - PAGE_OFFSET)) <= PAGE_TOLERANCE


def evaluate(rag, goldset, max_k):
    """Stellt jede Frage und merkt sich den Rang des ersten korrekten Treffers."""
    results = []
    for i, item in enumerate(goldset, start=1):
        query = item["question"]
        if USE_QUERY_INSTRUCTIONS:
            query = f"{query}\n\n{QUERY_ADDITIONAL_INSTRUCTIONS}"

        # Dieselbe Suche wie im Chat: Vektorsuche, ggf. Stichwortsuche und Reranking
        docs = [doc for doc, _ in rag.retrieve(query, k=max_k)]

        rank = None
        for position, doc in enumerate(docs, start=1):
            if is_hit(doc, item["source"], item["page"]):
                rank = position
                break

        results.append({**item, "rank": rank, "docs": docs})
        marker = f"Rang {rank}" if rank else "nicht gefunden"
        print(f"  [{i}/{len(goldset)}] {marker:<15} {item['question'][:60]}")

    return results


def judge(results, model):
    """Lässt ein Sprachmodell prüfen, welche der ersten JUDGE_K Treffer die Frage
    beantworten. Ergänzt je Frage judge_rank (erster passender Treffer) und judge_note."""
    import hashlib
    import json
    import re
    from langchain_openai import ChatOpenAI

    cache = json.loads(JUDGE_CACHE.read_text(encoding="utf-8")) if JUDGE_CACHE.exists() else {}
    llm = ChatOpenAI(model=model)
    for i, r in enumerate(results, start=1):
        docs = r["docs"][:JUDGE_K]
        excerpts = "\n\n".join(
            f"[{n}] ({Path(d.metadata.get('source', '')).name}, Seite {d.metadata.get('page')})\n"
            f"{' '.join(d.page_content.split())}" for n, d in enumerate(docs, start=1))
        prompt = JUDGE_PROMPT.format(n=len(docs), frage=r["question"], auszuege=excerpts)
        key = hashlib.sha256(f"{model}\n{prompt}".encode()).hexdigest()
        if key not in cache:
            reply = llm.invoke(prompt)
            text = reply.content if hasattr(reply, "content") else str(reply)
            if isinstance(text, list):          # manche Modelle liefern Inhaltsblöcke
                text = " ".join(part.get("text", "") if isinstance(part, dict) else str(part) for part in text)
            match = re.search(r"\{.*\}", text, re.S)
            try:
                verdict = json.loads(match.group(0)) if match else {}
            except json.JSONDecodeError:
                verdict = {}
            cache[key] = {"auszuege": sorted(int(n) for n in verdict.get("auszuege", []) if str(n).isdigit()),
                          "begruendung": verdict.get("begruendung", text[:200])}
            JUDGE_CACHE.write_text(json.dumps(cache, ensure_ascii=False, indent=1), encoding="utf-8")
        hits = [n for n in cache[key]["auszuege"] if 1 <= n <= len(docs)]
        r["judge_rank"] = hits[0] if hits else None
        r["judge_note"] = cache[key]["begruendung"]
        marker = f"Treffer {r['judge_rank']}" if hits else "keine Antwort"
        print(f"  [{i}/{len(results)}] {marker:<15} {r['question'][:60]}")


def judge_report(results, model):
    """Berichtsteil für die Prüfung durch das Sprachmodell."""
    n = len(results)
    found = [r for r in results if r.get("judge_rank")]
    page_hits = {id(r) for r in results if r["rank"] is not None and r["rank"] <= JUDGE_K}
    both = sum(1 for r in found if id(r) in page_hits)
    out = [
        "",
        f"Antwort im Kontext (Prüfung durch {model}, erste {JUDGE_K} Treffer)",
        f"  Antwort@{JUDGE_K:<11} {len(found) / n:.2f}   ({len(found)} von {n})",
        f"  MRR (Antwort){'':<6} {sum(1 / r['judge_rank'] for r in found) / n:.2f}",
        f"  Abgleich mit Seitenprüfung: beide ja {both}, nur Modell ja {len(found) - both}, "
        f"nur Seite ja {len(page_hits) - both}",
        "",
        "  Einzelne Fragen (Treffer = erster Auszug, der die Frage beantwortet; Seite = Rang laut Fundstelle)",
    ]
    for r in sorted(results, key=lambda r: (r.get("judge_rank") is None, r.get("judge_rank") or 0)):
        jr = r.get("judge_rank")
        pr = r["rank"] if r["rank"] is not None and r["rank"] <= JUDGE_K else None
        out.append(f"  Treffer {jr if jr else '–':>2} | Seite {pr if pr else '–':>2} | {r['question']}")
        out.append(f"      {r.get('judge_note', '')}")
    return out


def recall_at(results, k):
    """Anteil der Fragen, deren Fundstelle unter den ersten k Treffern liegt."""
    if not results:
        return 0.0
    found = sum(1 for r in results if r["rank"] is not None and r["rank"] <= k)
    return found / len(results)


def mrr(results):
    """Mean Reciprocal Rank: belohnt Treffer weit oben. 1.0 = immer Platz 1."""
    if not results:
        return 0.0
    return sum(1 / r["rank"] for r in results if r["rank"] is not None) / len(results)


def build_report(results, embedding_model):
    """Baut den Berichtstext, der auf dem Bildschirm und in der Datei landet."""
    line = "=" * 68
    out = [
        line,
        f"Retrieval-Auswertung — {datetime.now():%Y-%m-%d %H:%M}",
        line,
        "",
        "Konfiguration",
        f"  Vektorspeicher      {VECTORSTORE_PATH}",
        f"  Embedding-Modell    {embedding_model}",
        f"  Chunk-Größe         {INDEX_INFO.get('chunk_size', CHUNK_SIZE)} (Overlap {INDEX_INFO.get('chunk_overlap', CHUNK_OVERLAP)})",
        f"  Loader              {INDEX_INFO.get('pdf_loader', '?')}, Kopfzeilen {'an' if INDEX_INFO.get('chunk_headers') else 'aus'}",
        f"  Query-Zusatztext    {'an' if USE_QUERY_INSTRUCTIONS else 'aus'}",
        f"  Suche               {RETRIEVAL_SETUP}",
        f"  Seitentoleranz      ±{PAGE_TOLERANCE}",
        f"  Fragen              {len(results)}",
        "",
        "Ergebnis",
    ]

    for k in K_VALUES:
        value = recall_at(results, k)
        found = sum(1 for r in results if r["rank"] is not None and r["rank"] <= k)
        out.append(f"  Recall@{k:<12} {value:.2f}   ({found} von {len(results)})")
    out.append(f"  MRR{'':<16} {mrr(results):.2f}")

    corpora = sorted({r["corpus"] for r in results})
    if len(corpora) > 1:
        out += ["", "Nach Korpus"]
        smallest_k = min(K_VALUES)
        for corpus in corpora:
            subset = [r for r in results if r["corpus"] == corpus]
            wort = "Frage" if len(subset) == 1 else "Fragen"
            out.append(
                f"  {corpus:<12} Recall@{smallest_k} {recall_at(subset, smallest_k):.2f}"
                f"   MRR {mrr(subset):.2f}   ({len(subset)} {wort})"
            )

    out += ["", "Einzelne Fragen (Rang = Position des ersten korrekten Treffers)"]
    for r in sorted(results, key=lambda r: (r["rank"] is None, r["rank"] or 0)):
        rank = f"{r['rank']:>3}" if r["rank"] is not None else "  –"
        out.append(f"  Rang {rank}  {r['corpus']:<10} {r['question']}")

    nowhere = [r for r in results if r["rank"] is None]
    if nowhere:
        out += [
            "",
            f"{len(nowhere)} Frage(n) ohne Treffer unter den ersten {max(K_VALUES)}.",
            "Prüfen: Steht die Fundstelle wirklich dort? Stimmt der Dateiname?",
            "Ist das Dokument überhaupt im Index (rag.list_documents())?",
        ]

    if any("judge_rank" in r for r in results):
        out += judge_report(results, JUDGE_MODEL)

    out.append("")
    return "\n".join(out)


def main():
    global GOLDSET_PATH, VECTORSTORE_PATH, REPORT_PATH, USE_QUERY_INSTRUCTIONS
    parser = argparse.ArgumentParser(description="Retrieval-Auswertung")
    parser.add_argument("--goldset", type=Path, default=GOLDSET_PATH)
    parser.add_argument("--index", default=VECTORSTORE_PATH)
    parser.add_argument("--report", type=Path, default=REPORT_PATH)
    parser.add_argument("--mode", choices=["dense", "hybrid", "union"], default=None,
                        help="Suchmodus, Standard aus local_rag.RETRIEVAL_MODE")
    parser.add_argument("--no-rerank", action="store_true", help="Cross-Encoder abschalten")
    parser.add_argument("--judge", metavar="MODELL", default=None,
                        help="zusätzlich prüfen lassen, ob die Antwort in den ersten Treffern steht")
    parser.add_argument("--query-instructions", action="store_true", default=USE_QUERY_INSTRUCTIONS,
                        help="Zusatztext aus local_rag.QUERY_ADDITIONAL_INSTRUCTIONS anhängen")
    args = parser.parse_args()
    GOLDSET_PATH, VECTORSTORE_PATH, REPORT_PATH = args.goldset, args.index, args.report
    USE_QUERY_INSTRUCTIONS = args.query_instructions

    print(f"Lese Referenzfragen aus {GOLDSET_PATH} ...")
    goldset = read_goldset(GOLDSET_PATH)
    print(f"{len(goldset)} Frage(n) geladen.\n")

    # Lädt Embeddings und Vektorspeicher genau wie die Chat-Anwendung, damit
    # gemessen wird, was dort auch tatsächlich passiert. Ein LLM wird nur mit
    # --judge aufgerufen, zur Beurteilung der Treffer, nicht zum Antworten.
    options = {}
    if args.mode:
        options["retrieval_mode"] = args.mode
    if args.no_rerank:
        options["reranker_model"] = None
    rag = LocalRAG(vectorstore_path=VECTORSTORE_PATH, **options)
    if rag.vectorstore is None:
        sys.exit(
            f"\nKein Vektorspeicher unter {VECTORSTORE_PATH} gefunden.\n"
            "Erst den Index bauen (python add_pdfs.py), dann erneut auswerten."
        )

    print(f"\nStelle {len(goldset)} Frage(n) an den Vektorspeicher ...")
    global RETRIEVAL_SETUP, INDEX_INFO
    info_file = Path(VECTORSTORE_PATH) / "index_info.json"
    if info_file.exists():
        import json
        INDEX_INFO = json.loads(info_file.read_text(encoding="utf-8"))
    RETRIEVAL_SETUP = f"{rag.retrieval_mode}, Reranker {rag.reranker_model or 'aus'}"
    results = evaluate(rag, goldset, max_k=max(K_VALUES))
    if args.judge:
        global JUDGE_MODEL
        JUDGE_MODEL = args.judge
        print(f"\nLasse {args.judge} prüfen, ob die Antwort in den ersten {JUDGE_K} Treffern steht ...")
        judge(results, args.judge)

    report = build_report(results, rag.embeddings.model_name)
    print("\n" + report)

    REPORT_PATH.parent.mkdir(parents=True, exist_ok=True)
    REPORT_PATH.write_text(report, encoding="utf-8")
    print(f"Bericht gespeichert: {REPORT_PATH}")


if __name__ == "__main__":
    main()