from __future__ import annotations

import asyncio
from concurrent.futures import ThreadPoolExecutor
from threading import Lock

import pytest

from kyrodb import AsyncKyroDBClient, KyroDBClient
from kyrodb._generated import kyrodb_pb2 as pb2


class _SyncQueryStub:
    def __init__(self) -> None:
        self.calls = 0
        self._lock = Lock()

    def Query(
        self,
        request: pb2.QueryRequest,
        timeout: float | None = None,
        metadata: tuple[tuple[str, str], ...] | None = None,
    ) -> pb2.QueryResponse:
        _ = (timeout, metadata)
        with self._lock:
            self.calls += 1
        return pb2.QueryResponse(
            found=True,
            doc_id=request.doc_id,
            served_from=pb2.QueryResponse.HOT_TIER,
        )


class _AsyncSearchStub:
    def __init__(self) -> None:
        self.calls = 0
        self._lock = asyncio.Lock()

    async def Search(
        self,
        request: pb2.SearchRequest,
        timeout: float | None = None,
        metadata: tuple[tuple[str, str], ...] | None = None,
    ) -> pb2.SearchResponse:
        _ = (request, timeout, metadata)
        async with self._lock:
            self.calls += 1
        return pb2.SearchResponse(
            results=[pb2.SearchResult(doc_id=1, score=1.0)],
            total_found=1,
            search_latency_ms=0.1,
            search_path=pb2.SearchResponse.CACHE_HIT,
        )


def test_sync_client_query_is_thread_safe_under_fanout() -> None:
    with KyroDBClient(target="127.0.0.1:50051") as client:
        stub = _SyncQueryStub()
        client._stub = stub  # type: ignore[assignment]

        def call_query(doc_id: int) -> bool:
            return client.query(doc_id=doc_id, namespace="default").found

        with ThreadPoolExecutor(max_workers=8) as pool:
            results = list(pool.map(call_query, range(1, 101)))

        assert all(results)
        assert stub.calls == 100


@pytest.mark.asyncio
async def test_async_client_search_is_safe_under_concurrent_fanout() -> None:
    async with AsyncKyroDBClient(target="127.0.0.1:50051") as client:
        stub = _AsyncSearchStub()
        client._stub = stub  # type: ignore[assignment]

        async def call_search() -> int:
            response = await client.search(query_embedding=[0.1, 0.2], k=1, namespace="default")
            return response.total_found

        results = await asyncio.gather(*(call_search() for _ in range(100)))
        assert all(value == 1 for value in results)
        assert stub.calls == 100
