"""Performance benchmark for AiFastDb Python SDK.

T29.7: Measures throughput and latency for core operations.
Run with:  python benchmarks/bench_python_sdk.py

Reports results in a Markdown-friendly table for easy comparison
with Node.js SDK benchmarks.
"""

import os
import sys
import time
import tempfile
import statistics
from contextlib import contextmanager
from typing import Callable, List, Tuple

# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------

@contextmanager
def timer(label: str):
    """Context manager that prints elapsed time."""
    start = time.perf_counter()
    yield
    elapsed = time.perf_counter() - start
    print(f"  {label}: {elapsed:.4f}s")


def measure(fn: Callable, n: int = 1, warmup: int = 0) -> Tuple[float, float, float]:
    """Run fn() n times and return (total_sec, ops_per_sec, avg_ms)."""
    for _ in range(warmup):
        fn()
    start = time.perf_counter()
    for _ in range(n):
        fn()
    total = time.perf_counter() - start
    ops_per_sec = n / total if total > 0 else float("inf")
    avg_ms = (total / n) * 1000
    return total, ops_per_sec, avg_ms


def measure_latency(fn: Callable, n: int = 100) -> dict:
    """Run fn() n times and return latency stats."""
    times = []
    for _ in range(n):
        t0 = time.perf_counter()
        fn()
        times.append((time.perf_counter() - t0) * 1000)
    return {
        "min_ms": min(times),
        "max_ms": max(times),
        "avg_ms": statistics.mean(times),
        "p50_ms": statistics.median(times),
        "p99_ms": sorted(times)[int(len(times) * 0.99)],
    }


# ---------------------------------------------------------------------------
# Benchmarks
# ---------------------------------------------------------------------------

def bench_document_store(tmpdir: str):
    """DocumentStore write / read / query throughput."""
    from aifastdb import DocumentStore

    path = os.path.join(tmpdir, "bench_docs.jsonl")
    store = DocumentStore(path, enable_file_lock=False)

    N = 5000
    print(f"\n--- DocumentStore ({N} documents) ---")

    # Write throughput
    total, ops, avg = measure(
        lambda: None, 1  # placeholder
    )
    t0 = time.perf_counter()
    for i in range(N):
        store.put(f"doc-{i}", f"Content for document {i}", tags=["bench", f"batch-{i % 10}"])
    write_total = time.perf_counter() - t0
    print(f"  Write {N} docs:  {write_total:.4f}s  ({N / write_total:.0f} ops/s, {write_total / N * 1000:.3f} ms/op)")

    # Read throughput
    t0 = time.perf_counter()
    for i in range(N):
        store.get(f"doc-{i}")
    read_total = time.perf_counter() - t0
    print(f"  Read  {N} docs:  {read_total:.4f}s  ({N / read_total:.0f} ops/s, {read_total / N * 1000:.3f} ms/op)")

    # Query by tag
    lat = measure_latency(lambda: store.find_by_tag("bench"), n=100)
    print(f"  find_by_tag:  avg={lat['avg_ms']:.3f}ms  p50={lat['p50_ms']:.3f}ms  p99={lat['p99_ms']:.3f}ms")

    # Has (contains) check
    lat = measure_latency(lambda: store.has(f"doc-{N // 2}"), n=1000)
    print(f"  has():  avg={lat['avg_ms']:.4f}ms  p50={lat['p50_ms']:.4f}ms  p99={lat['p99_ms']:.4f}ms")

    store.flush()
    return {
        "doc_write_ops": N / write_total,
        "doc_read_ops": N / read_total,
    }


def bench_social_graph(tmpdir: str):
    """SocialGraphV2 person CRUD and path finding."""
    from aifastdb import SocialGraphV2

    path = os.path.join(tmpdir, "bench_social")
    sg = SocialGraphV2(path, wal_enabled=False)

    N_PERSONS = 1000
    print(f"\n--- SocialGraphV2 ({N_PERSONS} persons) ---")

    # Add persons
    t0 = time.perf_counter()
    person_ids = []
    for i in range(N_PERSONS):
        p = sg.add_person({"name": f"Person-{i}", "age": 20 + (i % 50)})
        person_ids.append(p["id"])
    add_total = time.perf_counter() - t0
    print(f"  Add {N_PERSONS} persons:  {add_total:.4f}s  ({N_PERSONS / add_total:.0f} ops/s)")

    # Add friendships (chain)
    t0 = time.perf_counter()
    n_friends = min(N_PERSONS - 1, 999)
    for i in range(n_friends):
        sg.add_friend(person_ids[i], person_ids[i + 1])
    friend_total = time.perf_counter() - t0
    print(f"  Add {n_friends} friendships:  {friend_total:.4f}s  ({n_friends / friend_total:.0f} ops/s)")

    # Get person
    lat = measure_latency(lambda: sg.get_person(person_ids[N_PERSONS // 2]), n=500)
    print(f"  get_person:  avg={lat['avg_ms']:.4f}ms  p50={lat['p50_ms']:.4f}ms  p99={lat['p99_ms']:.4f}ms")

    # Get friends
    lat = measure_latency(lambda: sg.get_friends(person_ids[N_PERSONS // 2]), n=200)
    print(f"  get_friends:  avg={lat['avg_ms']:.3f}ms  p50={lat['p50_ms']:.3f}ms  p99={lat['p99_ms']:.3f}ms")

    # Path finding (short path)
    lat = measure_latency(
        lambda: sg.find_path(person_ids[0], person_ids[min(5, N_PERSONS - 1)]),
        n=100,
    )
    print(f"  find_path (5 hops):  avg={lat['avg_ms']:.3f}ms  p50={lat['p50_ms']:.3f}ms  p99={lat['p99_ms']:.3f}ms")

    # List persons
    lat = measure_latency(lambda: sg.list_persons(), n=50)
    print(f"  list_persons ({N_PERSONS}):  avg={lat['avg_ms']:.2f}ms  p50={lat['p50_ms']:.2f}ms")

    sg.shutdown()
    return {
        "sg_add_person_ops": N_PERSONS / add_total,
        "sg_add_friend_ops": n_friends / friend_total,
    }


def bench_embedding(n: int = 10000):
    """EmbeddingEngine throughput."""
    from aifastdb import EmbeddingEngine

    engine = EmbeddingEngine()  # dummy 384d
    print(f"\n--- EmbeddingEngine (dummy 384d, {n} embeddings) ---")

    # Single embed
    t0 = time.perf_counter()
    for i in range(n):
        engine.embed(f"Text number {i}")
    embed_total = time.perf_counter() - t0
    print(f"  embed() × {n}:  {embed_total:.4f}s  ({n / embed_total:.0f} ops/s, {embed_total / n * 1000:.4f} ms/op)")

    # Batch embed
    batch_size = 100
    texts = [f"Text number {i}" for i in range(batch_size)]
    lat = measure_latency(lambda: engine.embed_batch(texts), n=100)
    print(f"  embed_batch({batch_size}):  avg={lat['avg_ms']:.3f}ms  ({batch_size / lat['avg_ms'] * 1000:.0f} texts/s)")

    # Numpy embed
    lat = measure_latency(lambda: engine.embed_numpy("hello"), n=1000)
    print(f"  embed_numpy():  avg={lat['avg_ms']:.4f}ms  p50={lat['p50_ms']:.4f}ms")

    # Cosine similarity
    v1, v2 = engine.embed("cat"), engine.embed("dog")
    lat = measure_latency(lambda: EmbeddingEngine.cosine_similarity(v1, v2), n=5000)
    print(f"  cosine_similarity:  avg={lat['avg_ms']:.4f}ms  ({1000 / lat['avg_ms']:.0f} ops/s)")

    return {
        "embed_ops": n / embed_total,
    }


def bench_federation(tmpdir: str):
    """Federation registration, query, and write."""
    from aifastdb import Federation

    print(f"\n--- Federation ---")

    fed = Federation()
    fed.register_graph("social", os.path.join(tmpdir, "fed_social"), wal_enabled=False)
    fed.register_document("docs", os.path.join(tmpdir, "fed_docs.jsonl"), enable_file_lock=False)

    N = 1000
    # Write
    t0 = time.perf_counter()
    for i in range(N):
        fed.write("docs", {
            "op": "CreateDocument",
            "id": f"fed-doc-{i}",
            "content": f"Federation document {i}",
            "doc_type": "note",
        })
    write_total = time.perf_counter() - t0
    print(f"  Federation write {N}:  {write_total:.4f}s  ({N / write_total:.0f} ops/s)")

    # Query
    lat = measure_latency(
        lambda: fed.query("docs", {"target": "document", "limit": 50}),
        n=100,
    )
    print(f"  Federation query:  avg={lat['avg_ms']:.3f}ms  p50={lat['p50_ms']:.3f}ms  p99={lat['p99_ms']:.3f}ms")

    fed.close()
    return {"fed_write_ops": N / write_total}


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    print("=" * 60)
    print("AiFastDb Python SDK — Performance Benchmark")
    print("=" * 60)
    print(f"Python: {sys.version}")

    try:
        import aifastdb
        print(f"aifastdb: {aifastdb.__version__}")
    except Exception:
        pass

    all_results = {}

    with tempfile.TemporaryDirectory(prefix="aifastdb_bench_") as tmpdir:
        all_results.update(bench_document_store(tmpdir))
        all_results.update(bench_social_graph(tmpdir))
        all_results.update(bench_embedding())
        all_results.update(bench_federation(tmpdir))

    # Summary table
    print("\n" + "=" * 60)
    print("Summary (ops/s)")
    print("=" * 60)
    print(f"| {'Operation':<35} | {'ops/s':>12} |")
    print(f"|{'-'*37}|{'-'*14}|")
    for key, val in all_results.items():
        label = key.replace("_ops", "").replace("_", " ").title()
        print(f"| {label:<35} | {val:>12,.0f} |")
    print()


if __name__ == "__main__":
    main()
