#!/usr/bin/env python3
"""
MCP-Zotero-Benchmark

Vergleicht den Zotero-MCP-Server (über SSH gegen lokale Zotero-API)
mit direkten Aufrufen der lokalen Zotero-HTTP-API und pyzotero.

Laufzeit wird absichtlich unter 2 Minuten gehalten. MCP-Tool-Aufrufe
werden mit Timeouts abgebrochen, damit einzelne hängende Operationen
den gesamten Benchmark nicht blockieren.

Start:
    python3 -m venv venv
    source venv/bin/activate
    pip install mcp pyzotero
    python mcp_zotero_benchmark.py
"""

import asyncio
import json
import os
import re
import sys
import time
import traceback
import urllib.request
import urllib.error
from datetime import datetime, timezone

# MCP SDK
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

# pyzotero
from pyzotero import zotero

CONFIG_PATH = os.environ.get("MCP_CONFIG", "~/.kimi-code/mcp.json")
LOCAL_API = os.environ.get("ZOTERO_LOCAL_API", "http://127.0.0.1:23119/api")
REPORT_PATH = os.environ.get("REPORT_PATH", "mcp_zotero_benchmark_report.md")
MCP_TOOL_TIMEOUT = int(os.environ.get("MCP_TOOL_TIMEOUT", "15"))

# Eine stabile Zotero-ID für wiederholbare Messungen.
DEFAULT_TEST_KEY = "YOUR_TEST_ITEM_KEY"


def load_config():
    with open(CONFIG_PATH, "r", encoding="utf-8") as f:
        cfg = json.load(f)
    z = cfg["mcpServers"]["zotero"]
    return z["command"], z["args"]


def parse_zotero_credentials_from_args(args_str):
    """Liest API-Key, Library-ID und Library-Type aus den SSH-Args des MCP-Servers."""
    key = re.search(r"ZOTERO_API_KEY='([^']+)'", args_str).group(1)
    lib_id = re.search(r"ZOTERO_LIBRARY_ID='([^']+)'", args_str).group(1)
    lib_type = re.search(r"ZOTERO_LIBRARY_TYPE='([^']+)'", args_str).group(1)
    return key, lib_id, lib_type


def now_iso():
    return datetime.now(timezone.utc).isoformat()


def timed(label, fn):
    """Führt fn() aus und gibt (Ergebnis, Dauer, Fehler) zurück."""
    t0 = time.perf_counter()
    try:
        result = fn()
        elapsed = time.perf_counter() - t0
        return result, elapsed, None
    except Exception as exc:
        elapsed = time.perf_counter() - t0
        return None, elapsed, exc


def print_result(label, elapsed, error=None, extra=""):
    if error is None:
        print(f"  {label}: {elapsed:.2f}s{extra}")
    else:
        print(f"  {label}: {elapsed:.2f}s – FEHLER: {error}{extra}")


async def mcp_call_with_timeout(session, name, arguments, timeout=MCP_TOOL_TIMEOUT):
    """Ruft ein MCP-Tool mit Timeout auf und gibt (Ergebnis, Dauer, Fehler) zurück."""
    t0 = time.perf_counter()
    try:
        result = await asyncio.wait_for(
            session.call_tool(name, arguments=arguments),
            timeout=timeout,
        )
        elapsed = time.perf_counter() - t0
        return result, elapsed, None
    except asyncio.TimeoutError:
        elapsed = time.perf_counter() - t0
        return None, elapsed, f"Timeout nach {timeout}s"
    except Exception as exc:
        elapsed = time.perf_counter() - t0
        return None, elapsed, str(exc)


async def benchmark_mcp(test_key: str):
    """Benchmark des MCP-Servers über SSH (stdio)."""
    command, args = load_config()
    server_params = StdioServerParameters(command=command, args=args, env=None)

    results = {}
    notes = []
    print("\n=== MCP-Server über SSH ===")

    try:
        t0 = time.perf_counter()
        async with stdio_client(server_params) as (read, write):
            async with ClientSession(read, write) as session:
                await session.initialize()
                results["init"] = time.perf_counter() - t0
                print(f"  Initialisierung (SSH+MCP): {results['init']:.2f}s")

                # list_tools
                t0 = time.perf_counter()
                tools = await session.list_tools()
                results["list_tools"] = time.perf_counter() - t0
                print(f"  list_tools ({len(tools.tools)} Tools): {results['list_tools']:.2f}s")

                # zotero_list_libraries – funktioniert zuverlässig
                result, elapsed, err = await mcp_call_with_timeout(
                    session, "zotero_list_libraries", {}
                )
                results["list_libraries"] = elapsed if err is None else None
                results["list_libraries_error"] = err
                print_result("list_libraries", elapsed, err)

                # zotero_get_recent – funktioniert, braucht aber etwas länger
                result, elapsed, err = await mcp_call_with_timeout(
                    session, "zotero_get_recent", {"limit": 5}
                )
                results["get_recent"] = elapsed if err is None else None
                results["get_recent_error"] = err
                print_result("get_recent(limit=5)", elapsed, err)

                # zotero_get_collections – funktioniert
                result, elapsed, err = await mcp_call_with_timeout(
                    session, "zotero_get_collections", {}
                )
                results["get_collections"] = elapsed if err is None else None
                results["get_collections_error"] = err
                print_result("get_collections", elapsed, err)

                # Item-Metadaten – in Standalone-Tests oft instabil
                result, elapsed, err = await mcp_call_with_timeout(
                    session, "zotero_get_item_metadata", {"item_key": test_key}
                )
                results["get_item_metadata"] = elapsed if err is None else None
                results["get_item_metadata_error"] = err
                print_result(f"get_item_metadata({test_key})", elapsed, err)

                # Kinder – in Standalone-Tests oft instabil
                result, elapsed, err = await mcp_call_with_timeout(
                    session, "zotero_get_item_children", {"item_key": test_key}
                )
                results["get_item_children"] = elapsed if err is None else None
                results["get_item_children_error"] = err
                print_result(f"get_item_children({test_key})", elapsed, err)

                # Volltext – in Standalone-Tests oft instabil
                result, elapsed, err = await mcp_call_with_timeout(
                    session, "zotero_get_item_fulltext", {"item_key": test_key}
                )
                results["get_item_fulltext"] = elapsed if err is None else None
                results["get_item_fulltext_error"] = err
                ft_len = 0
                if result and result.content:
                    ft_len = len(result.content[0].text)
                print_result(
                    f"get_item_fulltext({test_key})",
                    elapsed,
                    err,
                    f" ({ft_len:,} Zeichen)",
                )

                # Suche – in großen Bibliotheken langsam
                result, elapsed, err = await mcp_call_with_timeout(
                    session, "zotero_search_items", {"query": "AMSTAR-2", "limit": 5}
                )
                results["search_items"] = elapsed if err is None else None
                results["search_items_error"] = err
                print_result("search_items('AMSTAR-2', limit=5)", elapsed, err)

    except Exception as exc:
        print(f"  Schwerer Fehler im MCP-Benchmark: {exc}", file=sys.stderr)
        traceback.print_exc()
        results["fatal_error"] = str(exc)

    results["notes"] = notes
    return results


def find_pdf_attachment_key(children):
    """Sucht den wahrscheinlichsten Haupt-PDF-Attachment-Key.

    Bevorzugt Dateien, deren Name wie ein Artikel-PDF aussieht
    (nicht Supplement, nicht vom Nutzer erzeugte Zusammenfassungen).
    """
    pdfs = [
        c for c in children or []
        if c.get("data", {}).get("contentType") == "application/pdf"
    ]
    if not pdfs:
        return None

    def score(child):
        fn = (child.get("data", {}).get("filename") or "").lower()
        if any(x in fn for x in ["moesm", "suppl", "kimi", "beaver"]):
            return 1
        return 2

    pdfs.sort(key=score, reverse=True)
    return pdfs[0]["key"]


def benchmark_local_http_api(test_key: str, lib_type: str, lib_id: str):
    """Direkte Aufrufe der lokalen Zotero-HTTP-API (Port 23119)."""
    results = {}
    base_url = f"{LOCAL_API}/{lib_type}s/{lib_id}"
    headers = {"Accept": "application/json"}

    print("\n=== Lokale Zotero-HTTP-API (direkt) ===")

    # Kinder holen, um PDF-Attachment-Key zu finden
    children_req = urllib.request.Request(
        f"{base_url}/items/{test_key}/children", headers=headers
    )
    try:
        with urllib.request.urlopen(children_req, timeout=30) as resp:
            children = json.loads(resp.read())
    except Exception as e:
        print(f"  Kinder konnten nicht geladen werden: {e}", file=sys.stderr)
        children = []

    # Item-Metadaten
    def fetch_metadata():
        req = urllib.request.Request(
            f"{base_url}/items/{test_key}", headers=headers
        )
        with urllib.request.urlopen(req, timeout=30) as resp:
            return resp.read()

    body, elapsed, err = timed(f"GET items/{test_key}", fetch_metadata)
    results["item_metadata"] = elapsed
    results["item_metadata_error"] = str(err) if err else None
    print_result(f"GET items/{test_key}", elapsed, err)

    # Fulltext über den PDF-Attachment-Key (Parent-Key liefert 404)
    pdf_key = find_pdf_attachment_key(children)
    if pdf_key:
        def fetch_fulltext():
            req = urllib.request.Request(
                f"{base_url}/items/{pdf_key}/fulltext", headers=headers
            )
            with urllib.request.urlopen(req, timeout=30) as resp:
                data = json.loads(resp.read())
                return data.get("content", "").encode("utf-8")

        body, elapsed, err = timed(f"GET items/{pdf_key}/fulltext", fetch_fulltext)
        results["item_fulltext"] = elapsed
        if isinstance(err, urllib.error.HTTPError):
            err = f"HTTP {err.code}"
        results["item_fulltext_error"] = str(err) if err else None
        extra = f" ({len(body):,} Bytes)" if body else ""
        print_result(f"GET items/{pdf_key}/fulltext", elapsed, err, extra)
    else:
        results["item_fulltext"] = None
        results["item_fulltext_error"] = "kein PDF-Attachment"
        print("  GET .../fulltext: übersprungen (kein PDF-Attachment)")

    # Top-20 Items
    def fetch_top():
        req = urllib.request.Request(
            f"{base_url}/items?limit=20", headers=headers
        )
        with urllib.request.urlopen(req, timeout=30) as resp:
            return resp.read()

    body, elapsed, err = timed("GET items?limit=20", fetch_top)
    results["items_top20"] = elapsed
    results["items_top20_error"] = str(err) if err else None
    print_result("GET items?limit=20", elapsed, err)

    return results


def benchmark_pyzotero(test_key: str, lib_type: str, lib_id: str, api_key: str):
    """pyzotero-Benchmark: nutzt die lokale Zotero-HTTP-API (local=True)."""
    results = {}
    z = zotero.Zotero(lib_id, lib_type, api_key, local=True)

    print("\n=== pyzotero (lokale Zotero-API, local=True) ===")

    # Top 20 items
    items, elapsed, err = timed("items(limit=20)", lambda: z.items(limit=20))
    results["items_top20"] = elapsed
    results["items_top20_error"] = str(err) if err else None
    n_items = len(items) if items else 0
    print_result("items(limit=20)", elapsed, err, f" ({n_items} geliefert)")

    # Einzelnes Item
    item, elapsed, err = timed(f"item({test_key})", lambda: z.item(test_key))
    results["item"] = elapsed
    results["item_error"] = str(err) if err else None
    print_result(f"item({test_key})", elapsed, err)

    # Kinder
    children, elapsed, err = timed(
        f"children({test_key})", lambda: z.children(test_key)
    )
    results["children"] = elapsed
    results["children_error"] = str(err) if err else None
    n_children = len(children) if children else 0
    print_result(f"children({test_key})", elapsed, err, f" ({n_children} Kinder)")

    # Volltext über den ersten PDF-Attachment-Key
    pdf_key = find_pdf_attachment_key(children)
    if pdf_key:
        def fetch_fulltext():
            data = z.fulltext_item(pdf_key)
            return data.get("content", "")

        data, elapsed, err = timed(f"fulltext_item({pdf_key})", fetch_fulltext)
        results["item_fulltext"] = elapsed
        results["item_fulltext_error"] = str(err) if err else None
        extra = f" ({len(data):,} Zeichen)" if data else ""
        print_result(f"fulltext_item({pdf_key})", elapsed, err, extra)
    else:
        results["item_fulltext"] = None
        results["item_fulltext_error"] = "kein PDF-Attachment"
        print("  fulltext_item: übersprungen (kein PDF-Attachment)")

    return results


def fmt(v, err=None):
    if err is not None:
        return "—"
    if v is None:
        return "—"
    return f"{v:.2f}s"


def render_report(
    mcp_results, http_results, pyz_results, test_key, lib_id, lib_type, duration_total
):
    lines = [
        "# MCP-Zotero-Benchmark-Report",
        "",
        f"**Erstellt:** {now_iso()}  ",
        f"**Test-Item:** `{test_key}`  ",
        f"**Bibliothek:** {lib_type} {lib_id}  ",
        f"**Gesamtlaufzeit:** {duration_total:.1f}s",
        f"**MCP-Tool-Timeout:** {MCP_TOOL_TIMEOUT}s",
        "",
        "Alle Zeiten in Sekunden, gemessen auf dem Agenten-Rechner (Linux-VM). "
        "Der MCP-Server läuft auf dem Windows-Rechner (`YOUR_WINDOWS_HOST`) und "
        "wird über SSH+PowerShell gestartet.",
        "",
        "## Ergebnisse",
        "",
        "| Operation | MCP-Server (SSH) | Lokale HTTP-API | pyzotero (local=True) |",
        "|-----------|------------------|-----------------|-----------------------|",
    ]

    rows = [
        ("Initialisierung/Verbindung", mcp_results.get("init"), None, None, None, None, None),
        ("list_tools / —", mcp_results.get("list_tools"), None, None, None, None, None),
        ("list_libraries", mcp_results.get("list_libraries"), None, None, mcp_results.get("list_libraries_error"), None, None),
        ("get_recent(limit=5)", mcp_results.get("get_recent"), None, None, mcp_results.get("get_recent_error"), None, None),
        ("get_collections", mcp_results.get("get_collections"), None, None, mcp_results.get("get_collections_error"), None, None),
        ("Item-Metadaten", mcp_results.get("get_item_metadata"), http_results.get("item_metadata"), pyz_results.get("item"), mcp_results.get("get_item_metadata_error"), http_results.get("item_metadata_error"), pyz_results.get("item_error")),
        ("Kinder/Attachments", mcp_results.get("get_item_children"), None, pyz_results.get("children"), mcp_results.get("get_item_children_error"), None, pyz_results.get("children_error")),
        ("Volltext/PDF", mcp_results.get("get_item_fulltext"), http_results.get("item_fulltext"), pyz_results.get("item_fulltext"), mcp_results.get("get_item_fulltext_error"), http_results.get("item_fulltext_error"), pyz_results.get("item_fulltext_error")),
        ("Suche/Top-20 Items", mcp_results.get("search_items"), http_results.get("items_top20"), pyz_results.get("items_top20"), mcp_results.get("search_items_error"), http_results.get("items_top20_error"), pyz_results.get("items_top20_error")),
    ]

    for label, m, h, p, m_err, h_err, p_err in rows:
        lines.append(
            f"| {label} | {fmt(m, m_err)} | {fmt(h, h_err)} | {fmt(p, p_err)} |"
        )

    lines += [
        "",
        "## Beobachtungen",
        "",
        "- Der MCP-Server benötigt eine SSH-Verbindung und den Start von PowerShell; "
        "die Initialisierung dauert daher länger als ein direkter HTTP-Aufruf.",
        "- Sobald der Server läuft, liegen die meisten Tool-Aufrufe im Bereich "
        "weniger Sekunden. In manchen Läufen treten jedoch Timeouts oder "
        "Server-Abstürze (`AssertionError: Request already responded to`) auf; "
        "der Benchmark fängt diese mit Timeouts ab, damit die übrigen Messungen "
        "nicht blockiert werden.",
        "- Die lokale Zotero-HTTP-API arbeitet schneller und stabiler als der "
        "MCP-Server über SSH; für reine Lesezugriffe ist sie der effizienteste Weg.",
        "- `pyzotero` mit `local=True` greift auf dieselbe lokale Zotero-API zu "
        "wie die direkten HTTP-Aufrufe; Unterschiede in der Geschwindigkeit "
        "liegen dann am Overhead des Python-Clients.",
        "- PDF-Volltexte funktionieren über alle drei Wege, wenn man den "
        "Attachment-Key verwendet (lokale API und `pyzotero`) bzw. den Parent-Key "
        "(MCP-Server). Der MCP-Server löst den Attachment-Key intern auf.",
        "- Die lokale HTTP-API ist der schnellste Weg, liefert aber nur "
        "lesenden Zugriff und keine semantische Suche.",
        "",
        "## Konfiguration",
        "",
        "```json",
        json.dumps({"local_api": LOCAL_API, "config": CONFIG_PATH, "mcp_tool_timeout": MCP_TOOL_TIMEOUT}, indent=2),
        "```",
    ]

    errors = []
    if mcp_results.get("error"):
        errors.append(f"- MCP gesamt: `{mcp_results['error']}`")
    if mcp_results.get("fatal_error"):
        errors.append(f"- MCP fatal: `{mcp_results['fatal_error']}`")

    def collect_errors(prefix, results):
        for key, val in results.items():
            if key.endswith("_error") and val is not None:
                op = key[:-6]
                errors.append(f"- {prefix} `{op}`: `{val}`")

    collect_errors("MCP", mcp_results)
    collect_errors("Lokale HTTP-API", http_results)
    collect_errors("pyzotero", pyz_results)

    if errors:
        lines += ["", "## Fehler", ""] + errors

    return "\n".join(lines)


def main():
    command, args = load_config()
    args_str = " ".join(args)
    api_key, lib_id, lib_type = parse_zotero_credentials_from_args(args_str)

    test_key = DEFAULT_TEST_KEY

    print(f"Konfiguration: {CONFIG_PATH}")
    print(f"Lokale API:    {LOCAL_API}")
    print(f"Bibliothek:    {lib_type} {lib_id}")
    print(f"Test-Item:     {test_key}")
    print(f"Timeout:       {MCP_TOOL_TIMEOUT}s")
    print(f"Zeitstempel:   {now_iso()}")

    t_start = time.perf_counter()

    # MCP-Benchmark
    try:
        mcp_results = asyncio.run(benchmark_mcp(test_key))
    except Exception as e:
        print(f"Fehler im MCP-Benchmark: {e}", file=sys.stderr)
        traceback.print_exc()
        mcp_results = {"error": str(e)}

    # Lokale HTTP-API
    try:
        http_results = benchmark_local_http_api(test_key, lib_type, lib_id)
    except Exception as e:
        print(f"Fehler in lokaler HTTP-API: {e}", file=sys.stderr)
        traceback.print_exc()
        http_results = {"error": str(e)}

    # pyzotero (local=True)
    try:
        pyz_results = benchmark_pyzotero(test_key, lib_type, lib_id, api_key)
    except Exception as e:
        print(f"Fehler in pyzotero: {e}", file=sys.stderr)
        traceback.print_exc()
        pyz_results = {"error": str(e)}

    duration_total = time.perf_counter() - t_start

    report = render_report(
        mcp_results, http_results, pyz_results, test_key, lib_id, lib_type, duration_total
    )

    with open(REPORT_PATH, "w", encoding="utf-8") as f:
        f.write(report)

    print(f"\nGesamtlaufzeit: {duration_total:.1f}s")
    print(f"Report geschrieben nach: {REPORT_PATH}")


if __name__ == "__main__":
    main()