# @invar:allow shell_result: Benchmark functions output directly to stdout, no Result pattern needed
# @invar:allow shell_pure_logic: Benchmark timing functions use time.perf_counter() and random
"""
Performance benchmarks for Lattice search operations.

Benchmarks:
1. FTS5 full-text search (BM25)
2. Vector similarity search
3. RRF (Reciprocal Rank Fusion)
4. Hybrid search (FTS5 + Vector + RRF)

Run with: python benchmarks/search.py
"""

from __future__ import annotations

import sqlite3
import statistics
import tempfile
import time
from pathlib import Path

from returns.result import Success

from lattice.core.search import (
    calculate_rrf_score,
    fuse_results,
)
from lattice.core.types.search import FtsResult, VecResult
from lattice.shell.schema import create_store
from lattice.shell.store import insert_log, search_fts
from lattice.core.types.enums import Role


# ==============================================================================
# Benchmark Utilities
# ==============================================================================


def benchmark_result(
    name: str,
    times: list[float],
    iterations: int,
) -> dict[str, float | int | str]:
    """Calculate benchmark statistics from timing data.

    Args:
        name: Benchmark name.
        times: List of timing measurements in seconds.
        iterations: Number of iterations run.

    Returns:
        Dictionary with mean, median, p95, min, max, and iterations.
    """
    sorted_times = sorted(times)
    p95_index = int(len(sorted_times) * 0.95)

    return {
        "name": name,
        "mean_ms": statistics.mean(times) * 1000,
        "median_ms": statistics.median(times) * 1000,
        "p95_ms": sorted_times[p95_index] * 1000 if sorted_times else 0.0,
        "min_ms": min(times) * 1000 if times else 0.0,
        "max_ms": max(times) * 1000 if times else 0.0,
        "iterations": iterations,
    }


def print_results(results: list[dict[str, float | int | str]]) -> None:
    """Print benchmark results in a formatted table.

    Args:
        results: List of benchmark result dictionaries.
    """
    print("\n" + "=" * 80)
    print("BENCHMARK RESULTS")
    print("=" * 80)
    print(
        f"{'Benchmark':<40} {'Mean (ms)':>10} {'Median (ms)':>12} {'P95 (ms)':>10} {'Iter':>6}"
    )
    print("-" * 80)
    for r in results:
        print(
            f"{r['name']:<40} {r['mean_ms']:>10.3f} {r['median_ms']:>12.3f} "
            f"{r['p95_ms']:>10.3f} {r['iterations']:>6}"
        )
    print("=" * 80 + "\n")


# ==============================================================================
# Test Data Setup
# ==============================================================================


def setup_test_database(
    num_logs: int = 1000,
) -> tuple[sqlite3.Connection, Path]:
    """Create a test database with sample log entries.

    Args:
        num_logs: Number of log entries to create.

    Returns:
        Tuple of (connection, db_path).
    """
    tmpdir = tempfile.mkdtemp()
    db_path = Path(tmpdir) / "bench_store.db"

    result = create_store(db_path)
    conn = result.unwrap()

    # Insert test logs with varied content for FTS5 search
    test_contents = [
        "The quick brown fox jumps over the lazy dog",
        "Machine learning models require training data",
        "Python is a versatile programming language",
        "Database indexing improves query performance",
        "Search algorithms are fundamental to information retrieval",
        "Natural language processing enables text analysis",
        "Vector embeddings capture semantic meaning",
        "BM25 is a popular ranking function for text search",
        "Reciprocal rank fusion combines multiple rankings",
        "SQLite FTS5 provides full-text search capabilities",
    ]

    for i in range(num_logs):
        content = test_contents[i % len(test_contents)]
        insert_log(
            conn,
            session_id=f"session-{i // 10}",
            role=Role.USER if i % 2 == 0 else Role.ASSISTANT,
            content=f"{content} - entry {i}",
        )

    return conn, db_path


def generate_test_vectors(count: int, dimensions: int = 384) -> list[list[float]]:
    """Generate test embedding vectors.

    Args:
        count: Number of vectors to generate.
        dimensions: Vector dimensions (default 384 for common models).

    Returns:
        List of normalized embedding vectors.
    """
    import random

    vectors = []
    for i in range(count):
        # Create deterministic but varied vectors
        random.seed(42 + i)
        vec = [random.gauss(0, 1) for _ in range(dimensions)]
        # Normalize
        magnitude = sum(x * x for x in vec) ** 0.5
        vectors.append([x / magnitude for x in vec])
    return vectors


# ==============================================================================
# Core Benchmarks (Pure Python, No I/O)
# ==============================================================================


def benchmark_rrf_score_calculation(
    iterations: int = 10000,
) -> dict[str, float | int | str]:
    """Benchmark pure RRF score calculation (no I/O).

    Tests calculate_rrf_score() from core.search module.
    """
    times = []

    # Warmup
    for _ in range(100):
        calculate_rrf_score(1, 1, 60)

    # Benchmark
    for _ in range(iterations):
        start = time.perf_counter()
        calculate_rrf_score(1, 1, 60)
        times.append(time.perf_counter() - start)

    return benchmark_result("core.calculate_rrf_score", times, iterations)


def benchmark_rrf_fusion(iterations: int = 1000) -> dict[str, float | int | str]:
    """Benchmark RRF fusion of result lists (no I/O).

    Tests fuse_results() from core.search module.
    """
    # Create test data with overlapping results
    bm25_results = [FtsResult(rowid=i, rank=-0.1 * i) for i in range(1, 101)]
    vec_results = [VecResult(rowid=i, distance=0.01 * i) for i in range(1, 101)]

    times = []

    # Warmup
    for _ in range(10):
        fuse_results(bm25_results, vec_results)

    # Benchmark
    for _ in range(iterations):
        start = time.perf_counter()
        fuse_results(bm25_results, vec_results)
        times.append(time.perf_counter() - start)

    return benchmark_result("core.fuse_results (100+100)", times, iterations)


def benchmark_rrf_fusion_large(iterations: int = 100) -> dict[str, float | int | str]:
    """Benchmark RRF fusion with larger result sets (no I/O)."""
    # Create larger test data
    bm25_results = [FtsResult(rowid=i, rank=-0.1 * i) for i in range(1, 1001)]
    vec_results = [VecResult(rowid=i, distance=0.01 * i) for i in range(1, 1001)]

    times = []

    # Warmup
    for _ in range(5):
        fuse_results(bm25_results, vec_results)

    # Benchmark
    for _ in range(iterations):
        start = time.perf_counter()
        fuse_results(bm25_results, vec_results)
        times.append(time.perf_counter() - start)

    return benchmark_result("core.fuse_results (1000+1000)", times, iterations)


# ==============================================================================
# Shell Benchmarks (Database I/O)
# ==============================================================================


def benchmark_fts5_search(
    conn: sqlite3.Connection,
    query: str = "python",
    iterations: int = 100,
) -> dict[str, float | int | str]:
    """Benchmark FTS5 full-text search with BM25 ranking.

    Tests search_fts() from shell.store module.
    """
    times = []

    # Warmup
    for _ in range(10):
        search_fts(conn, query, limit=10)

    # Benchmark
    for _ in range(iterations):
        start = time.perf_counter()
        result = search_fts(conn, query, limit=10)
        times.append(time.perf_counter() - start)

    return benchmark_result("shell.search_fts (limit=10)", times, iterations)


def benchmark_fts5_search_large_limit(
    conn: sqlite3.Connection,
    query: str = "search",
    iterations: int = 100,
) -> dict[str, float | int | str]:
    """Benchmark FTS5 search with larger result set."""
    times = []

    # Warmup
    for _ in range(10):
        search_fts(conn, query, limit=100)

    # Benchmark
    for _ in range(iterations):
        start = time.perf_counter()
        result = search_fts(conn, query, limit=100)
        times.append(time.perf_counter() - start)

    return benchmark_result("shell.search_fts (limit=100)", times, iterations)


def benchmark_fts5_search_complex_query(
    conn: sqlite3.Connection,
    iterations: int = 100,
) -> dict[str, float | int | str]:
    """Benchmark FTS5 with complex query (AND/OR)."""
    complex_query = "python AND database"
    times = []

    # Warmup
    for _ in range(10):
        search_fts(conn, complex_query, limit=10)

    # Benchmark
    for _ in range(iterations):
        start = time.perf_counter()
        result = search_fts(conn, complex_query, limit=10)
        times.append(time.perf_counter() - start)

    return benchmark_result("shell.search_fts (complex AND)", times, iterations)


# ==============================================================================
# Hybrid Search Benchmark
# ==============================================================================


def benchmark_hybrid_search_simulated(
    conn: sqlite3.Connection,
    iterations: int = 100,
) -> dict[str, float | int | str]:
    """Benchmark simulated hybrid search (FTS5 + RRF fusion).

    Simulates the hybrid search workflow:
    1. FTS5 search for BM25 results
    2. Simulated vector search results
    3. RRF fusion
    """
    query = "search"
    times = []

    # Generate simulated vector results (since sqlite-vec may not be configured)
    vec_results = [VecResult(rowid=i, distance=0.01 * i) for i in range(1, 51)]

    # Warmup
    for _ in range(10):
        fts_result = search_fts(conn, query, limit=50)
        if isinstance(fts_result, Success):
            fused = fuse_results(fts_result.unwrap(), vec_results)

    # Benchmark
    for _ in range(iterations):
        start = time.perf_counter()
        fts_result = search_fts(conn, query, limit=50)
        if isinstance(fts_result, Success):
            fused = fuse_results(fts_result.unwrap(), vec_results)
        times.append(time.perf_counter() - start)

    return benchmark_result("shell.hybrid_search_simulated", times, iterations)


# ==============================================================================
# Main Entry Point
# ==============================================================================


def run_all_benchmarks(
    iterations_small: int = 1000,
    iterations_db: int = 100,
    num_logs: int = 1000,
) -> list[dict[str, float | int | str]]:
    """Run all benchmarks and return results.

    Args:
        iterations_small: Iterations for pure Python benchmarks.
        iterations_db: Iterations for database benchmarks.
        num_logs: Number of test log entries.

    Returns:
        List of benchmark result dictionaries.
    """
    print(
        f"\nRunning benchmarks (small={iterations_small}, db={iterations_db}, logs={num_logs})...\n"
    )

    results: list[dict[str, float | int | str]] = []

    # Core benchmarks (pure Python, fast)
    print("  [1/6] RRF score calculation...")
    results.append(benchmark_rrf_score_calculation(iterations_small))

    print("  [2/6] RRF fusion (100 results)...")
    results.append(benchmark_rrf_fusion(iterations_small // 10))

    print("  [3/6] RRF fusion (1000 results)...")
    results.append(benchmark_rrf_fusion_large(iterations_small // 10))

    # Database benchmarks
    print(f"  [4/6] Setting up test database ({num_logs} logs)...")
    conn, db_path = setup_test_database(num_logs)

    try:
        print("  [5/6] FTS5 search benchmarks...")
        results.append(benchmark_fts5_search(conn, iterations=iterations_db))
        results.append(
            benchmark_fts5_search_large_limit(conn, iterations=iterations_db)
        )
        results.append(
            benchmark_fts5_search_complex_query(conn, iterations=iterations_db)
        )

        print("  [6/6] Hybrid search benchmark...")
        results.append(
            benchmark_hybrid_search_simulated(conn, iterations=iterations_db)
        )
    finally:
        conn.close()
        # Cleanup temp directory
        import shutil

        shutil.rmtree(db_path.parent, ignore_errors=True)

    return results


def main() -> None:
    """Run benchmarks and print results."""
    print("=" * 80)
    print("LATTICE SEARCH PERFORMANCE BENCHMARKS")
    print("=" * 80)

    results = run_all_benchmarks()
    print_results(results)

    # Summary
    print("SUMMARY")
    print("-" * 40)
    for r in results:
        if "core." in str(r["name"]):
            print(f"  {r['name']}: {r['mean_ms']:.4f} ms (pure Python)")
        else:
            print(f"  {r['name']}: {r['mean_ms']:.4f} ms (with I/O)")


if __name__ == "__main__":
    main()
