"""RAG retrieval benchmark: p95 latency of get_conversation_context.

Standalone dev tool (do not import into the package). Seeds temporary
SQLite databases with 1k and 5k user/response rows using deterministic
fake embeddings (the embedding HTTP call is monkeypatched, so no network
is needed), then measures p95 of ChatDatabase.get_conversation_context
over repeated queries.

Run from the repo root:

    uv run python scripts/bench_rag.py

Prints p95@1k and p95@5k in milliseconds.
"""

from __future__ import annotations

import hashlib
import sys
import tempfile
import time
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import vibe_bot.db.messages as db_messages
from vibe_bot import llm_client
from vibe_bot.database import ChatDatabase

EMBEDDING_DIM = 256
NUM_CLUSTERS = 16
TOPICS = [
    "sailing",
    "baking",
    "gardening",
    "astronomy",
    "chess",
    "pottery",
    "mountaineering",
    "photography",
    "brewing",
    "carpentry",
    "weaving",
    "falconry",
    "cartography",
    "metallurgy",
    "botany",
    "masonry",
]
SIZES = (1000, 5000)
WARMUP_QUERIES = 20
MEASURED_QUERIES = 200
USER_ID = "bench-user"


def fake_embedding(text: str, *, model: str, url: str, api_key: str) -> list[float]:
    """Deterministic cluster-structured fake embedding (no network)."""
    topic = text.split(maxsplit=1)[0].lower()
    seed = int.from_bytes(hashlib.sha256(topic.encode("utf-8")).digest()[:8], "big")
    rng = np.random.default_rng(seed)
    center = np.zeros(EMBEDDING_DIM, dtype=np.float32)
    span = EMBEDDING_DIM // NUM_CLUSTERS
    cluster = seed % NUM_CLUSTERS
    center[cluster * span : (cluster + 1) * span] = 1.0
    noise = rng.standard_normal(EMBEDDING_DIM, dtype=np.float32)
    noise /= np.linalg.norm(noise)
    vector = center + 0.3 * noise
    return [float(x) for x in vector]


def seed_db(db_path: str, n_rows: int) -> list[str]:
    """Seed n_rows user/response pairs; return the user contents as queries."""
    db = ChatDatabase(db_path=db_path)
    queries: list[str] = []
    for i in range(n_rows):
        topic = TOPICS[i % NUM_CLUSTERS]
        content = f"{topic} question number {i}"
        queries.append(content)
        db.add_message(
            message_id=f"bench-{i}",
            user_id=USER_ID,
            username="bench",
            content=content,
        )
        db.add_message(
            message_id=f"bench-{i}_response",
            user_id="bench-bot",
            username="bench-bot",
            content=f"response {i}",
            role="assistant",
            embed=False,
        )
    return queries


def measure(db_path: str, queries: list[str]) -> float:
    """Return p95 in ms of get_conversation_context over repeated queries."""
    db = ChatDatabase(db_path=db_path)
    for query in queries[:WARMUP_QUERIES]:
        db.get_conversation_context(USER_ID, query)

    latencies_ms: list[float] = []
    for query in queries[:MEASURED_QUERIES]:
        start = time.perf_counter()
        db.get_conversation_context(USER_ID, query)
        latencies_ms.append((time.perf_counter() - start) * 1000)
    return float(np.percentile(latencies_ms, 95))


def main() -> None:
    llm_client.embedding = fake_embedding  # type: ignore[assignment]
    db_messages.MAX_HISTORY_MESSAGES = 10**9

    results: dict[int, float] = {}
    for size in SIZES:
        with tempfile.TemporaryDirectory(prefix="bench_rag_") as tmp:
            db_path = str(Path(tmp) / f"bench_{size}.db")
            queries = seed_db(db_path, size)
            results[size] = measure(db_path, queries)

    p95_1k = results[1000]
    p95_5k = results[5000]
    passed = p95_1k <= 10.0 and p95_5k <= 75.0
    print(f"p95@1k  = {p95_1k:.2f} ms")
    print(f"p95@5k  = {p95_5k:.2f} ms")
    print(f"ratio   = {p95_5k / p95_1k:.2f}x (informational)")
    print(f"target: p95@1k <= 10ms, p95@5k <= 75ms -> {'PASS' if passed else 'FAIL'}")


if __name__ == "__main__":
    main()