from __future__ import annotations

import asyncio
from collections.abc import AsyncIterable, Iterable
from dataclasses import dataclass

import grpc
import pytest

from kyrodb import AsyncKyroDBClient
from kyrodb._generated import kyrodb_pb2 as pb2
from kyrodb.errors import CircuitOpenError, DeadlineExceededError, ServiceUnavailableError
from kyrodb.models import BatchDeleteResult, BulkQueryResult, DeleteResult, InsertItem
from kyrodb.retry import CircuitBreakerPolicy, RetryPolicy

try:
    import numpy as np
except Exception:  # pragma: no cover
    np = None  # type: ignore[assignment]


@dataclass
class _CapturedCall:
    request: pb2.InsertRequest | None = None
    bulk_requests: list[pb2.InsertRequest] | None = None
    delete_request: pb2.DeleteRequest | None = None
    bulk_query_request: pb2.BulkQueryRequest | None = None
    batch_delete_request: pb2.BatchDeleteRequest | None = None
    metadata: tuple[tuple[str, str], ...] | None = None
    timeout: float | None = None
    bulk_query_calls: int = 0
    batch_delete_calls: int = 0


class _FakeStub:
    def __init__(self, captured: _CapturedCall) -> None:
        self._captured = captured

    async def Insert(
        self,
        request: pb2.InsertRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.InsertResponse:
        self._captured.request = request
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.InsertResponse(
            success=True,
            inserted_at=123,
            total_inserted=1,
            total_failed=0,
            tier=pb2.InsertResponse.HOT_TIER,
        )

    async def BulkInsert(
        self,
        requests: AsyncIterable[pb2.InsertRequest],
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.InsertResponse:
        self._captured.bulk_requests = []
        async for request in requests:
            self._captured.bulk_requests.append(request)
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.InsertResponse(
            success=True,
            inserted_at=123,
            total_inserted=len(self._captured.bulk_requests),
            total_failed=0,
            tier=pb2.InsertResponse.HOT_TIER,
        )

    async def BulkQuery(
        self,
        request: pb2.BulkQueryRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.BulkQueryResponse:
        self._captured.bulk_query_request = request
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        self._captured.bulk_query_calls += 1
        return pb2.BulkQueryResponse(total_found=0, total_requested=len(request.doc_ids))

    async def Delete(
        self,
        request: pb2.DeleteRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.DeleteResponse:
        self._captured.delete_request = request
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.DeleteResponse(success=True, existed=request.doc_id == 42)

    async def BatchDelete(
        self,
        request: pb2.BatchDeleteRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.BatchDeleteResponse:
        self._captured.batch_delete_request = request
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        self._captured.batch_delete_calls += 1
        return pb2.BatchDeleteResponse(success=True, deleted_count=len(request.ids.doc_ids))


class _RetryableRpcError(grpc.RpcError):
    def __init__(self, code: grpc.StatusCode) -> None:
        super().__init__()
        self._code = code

    def code(self) -> grpc.StatusCode:
        return self._code


class _RetryQueryStub:
    def __init__(self) -> None:
        self.attempts = 0
        self.call_metadata: list[tuple[tuple[str, str], ...]] = []

    async def Query(
        self,
        request: pb2.QueryRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.QueryResponse:
        _ = (request, timeout)
        self.attempts += 1
        self.call_metadata.append(tuple(metadata or ()))
        if self.attempts == 1:
            raise _RetryableRpcError(grpc.StatusCode.UNAVAILABLE)
        return pb2.QueryResponse(found=True, doc_id=1, served_from=pb2.QueryResponse.HOT_TIER)


class _AlwaysUnavailableQueryStub:
    def __init__(self) -> None:
        self.attempts = 0

    async def Query(
        self,
        request: pb2.QueryRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.QueryResponse:
        _ = (request, timeout, metadata)
        self.attempts += 1
        raise _RetryableRpcError(grpc.StatusCode.UNAVAILABLE)


@pytest.mark.asyncio
async def test_async_insert_supports_metadata() -> None:
    client = AsyncKyroDBClient(target="127.0.0.1:50051", api_key="k_test")
    captured = _CapturedCall()
    client._stub = _FakeStub(captured)  # type: ignore[assignment]

    ack = await client.insert(
        doc_id=42,
        embedding=[0.1, 0.2, 0.3],
        metadata={"source": "unit-test"},
        timeout_s=1.5,
    )

    assert ack.success is True
    assert captured.request is not None
    assert dict(captured.request.metadata) == {"source": "unit-test"}
    assert captured.timeout == 1.5
    assert captured.metadata is not None
    assert ("x-api-key", "k_test") in captured.metadata

    await client.close()


@pytest.mark.asyncio
async def test_async_bulk_insert_consumes_requests() -> None:
    client = AsyncKyroDBClient(target="127.0.0.1:50051", api_key="k_test")
    captured = _CapturedCall()
    client._stub = _FakeStub(captured)  # type: ignore[assignment]

    ack = await client.bulk_insert(
        [
            InsertItem.from_parts(doc_id=1, embedding=[0.1, 0.2], metadata={"a": "1"}),
            InsertItem.from_parts(doc_id=2, embedding=[0.3, 0.4], metadata={"b": "2"}),
        ],
        timeout_s=2.0,
    )

    assert ack.success is True
    assert captured.bulk_requests is not None
    assert len(captured.bulk_requests) == 2
    assert dict(captured.bulk_requests[0].metadata) == {"a": "1"}
    assert captured.timeout == 2.0
    assert captured.metadata is not None
    assert ("x-api-key", "k_test") in captured.metadata

    await client.close()


@pytest.mark.asyncio
async def test_async_bulk_query_and_batch_delete_validate_payloads() -> None:
    client = AsyncKyroDBClient(target="127.0.0.1:50051", api_key="k_test")
    captured = _CapturedCall()
    client._stub = _FakeStub(captured)  # type: ignore[assignment]

    bulk_query_response = await client.bulk_query(doc_ids=[10, 11], include_embeddings=True)
    assert isinstance(bulk_query_response, BulkQueryResult)
    assert bulk_query_response.total_requested == 2
    assert captured.bulk_query_request is not None
    assert list(captured.bulk_query_request.doc_ids) == [10, 11]
    assert captured.bulk_query_request.include_embeddings is True

    delete_response = await client.batch_delete_ids(doc_ids=[10, 11, 12])
    assert isinstance(delete_response, BatchDeleteResult)
    assert delete_response.success is True
    assert captured.batch_delete_request is not None
    assert list(captured.batch_delete_request.ids.doc_ids) == [10, 11, 12]

    await client.close()


@pytest.mark.asyncio
async def test_async_delete_returns_typed_result() -> None:
    client = AsyncKyroDBClient(target="127.0.0.1:50051")
    captured = _CapturedCall()
    client._stub = _FakeStub(captured)  # type: ignore[assignment]

    delete = await client.delete(doc_id=42, namespace="default", timeout_s=0.75)

    assert isinstance(delete, DeleteResult)
    assert delete.success is True
    assert delete.existed is True
    assert captured.delete_request is not None
    assert captured.delete_request.doc_id == 42
    assert captured.delete_request.namespace == "default"
    assert captured.timeout == 0.75
    await client.close()


@pytest.mark.asyncio
async def test_async_wait_for_ready_times_out() -> None:
    client = AsyncKyroDBClient(target="127.0.0.1:65535")
    with pytest.raises(DeadlineExceededError, match="WaitForReady"):
        await client.wait_for_ready(timeout_s=0.01)
    await client.close()


@pytest.mark.asyncio
async def test_async_wait_for_ready_maps_asyncio_timeout(monkeypatch: pytest.MonkeyPatch) -> None:
    class _DummyChannel:
        async def channel_ready(self) -> None:
            return None

        async def close(self) -> None:
            return None

    class _DummyStub:
        def __init__(self, _channel: object) -> None:
            return None

    async def fake_wait_for(awaitable: object, timeout: float | None = None) -> None:
        _ = timeout
        close = getattr(awaitable, "close", None)
        if callable(close):
            close()
        raise asyncio.TimeoutError

    monkeypatch.setattr(
        "kyrodb.async_client.grpc.aio.insecure_channel",
        lambda *args, **kwargs: _DummyChannel(),
    )
    monkeypatch.setattr("kyrodb.async_client.pb2_grpc.KyroDBServiceStub", _DummyStub)
    monkeypatch.setattr("kyrodb.async_client.asyncio.wait_for", fake_wait_for)

    client = AsyncKyroDBClient(target="127.0.0.1:50051")
    with pytest.raises(DeadlineExceededError, match="WaitForReady"):
        await client.wait_for_ready(timeout_s=0.5)
    await client.close()


def test_async_client_rejects_non_positive_default_timeout() -> None:
    with pytest.raises(ValueError, match="default_timeout_s"):
        AsyncKyroDBClient(target="127.0.0.1:50051", default_timeout_s=0)


def test_async_client_rejects_non_positive_max_unary_batch_size() -> None:
    with pytest.raises(ValueError, match="max_unary_batch_size"):
        AsyncKyroDBClient(target="127.0.0.1:50051", max_unary_batch_size=0)


def test_async_client_rejects_blank_lb_policy_name() -> None:
    with pytest.raises(ValueError, match="lb_policy_name"):
        AsyncKyroDBClient(target="127.0.0.1:50051", lb_policy_name="  ")


@pytest.mark.asyncio
async def test_async_api_key_rotation_with_set_api_key() -> None:
    client = AsyncKyroDBClient(target="127.0.0.1:50051", api_key="key_v1")
    captured = _CapturedCall()
    client._stub = _FakeStub(captured)  # type: ignore[assignment]

    await client.insert(doc_id=1, embedding=[0.1, 0.2])
    assert captured.metadata is not None
    assert ("x-api-key", "key_v1") in captured.metadata

    client.set_api_key("key_v2")
    await client.insert(doc_id=2, embedding=[0.1, 0.2])
    assert captured.metadata is not None
    assert ("x-api-key", "key_v2") in captured.metadata

    await client.close()


@pytest.mark.asyncio
async def test_async_api_key_provider_is_used_per_call() -> None:
    key_state = {"value": "key_v1"}

    def provider() -> str:
        return key_state["value"]

    client = AsyncKyroDBClient(target="127.0.0.1:50051")
    captured = _CapturedCall()
    client._stub = _FakeStub(captured)  # type: ignore[assignment]
    client.set_api_key_provider(provider)

    await client.insert(doc_id=1, embedding=[0.1, 0.2])
    assert captured.metadata is not None
    assert ("x-api-key", "key_v1") in captured.metadata

    key_state["value"] = "key_v2"
    await client.insert(doc_id=2, embedding=[0.1, 0.2])
    assert captured.metadata is not None
    assert ("x-api-key", "key_v2") in captured.metadata

    await client.close()


@pytest.mark.asyncio
async def test_async_api_key_provider_async_is_used_per_call() -> None:
    key_state = {"value": "key_v1"}

    async def provider() -> str:
        return key_state["value"]

    client = AsyncKyroDBClient(target="127.0.0.1:50051")
    captured = _CapturedCall()
    client._stub = _FakeStub(captured)  # type: ignore[assignment]
    await client.set_api_key_provider_async(provider)

    await client.insert(doc_id=1, embedding=[0.1, 0.2])
    assert captured.metadata is not None
    assert ("x-api-key", "key_v1") in captured.metadata

    key_state["value"] = "key_v2"
    await client.insert(doc_id=2, embedding=[0.1, 0.2])
    assert captured.metadata is not None
    assert ("x-api-key", "key_v2") in captured.metadata

    await client.close()


@pytest.mark.asyncio
async def test_async_call_timeout_validates_positive_values() -> None:
    client = AsyncKyroDBClient(target="127.0.0.1:50051")
    captured = _CapturedCall()
    client._stub = _FakeStub(captured)  # type: ignore[assignment]
    with pytest.raises(ValueError, match="timeout_s must be > 0"):
        await client.insert(doc_id=1, embedding=[0.1, 0.2], timeout_s=0)
    await client.close()


@pytest.mark.asyncio
async def test_async_api_key_provider_requires_callable() -> None:
    client = AsyncKyroDBClient(target="127.0.0.1:50051")
    with pytest.raises(TypeError, match="provider must be callable"):
        client.set_api_key_provider("not-callable")  # type: ignore[arg-type]
    await client.close()


@pytest.mark.asyncio
async def test_async_retry_preserves_call_metadata_for_generators() -> None:
    client = AsyncKyroDBClient(
        target="127.0.0.1:50051",
        retry_policy=RetryPolicy(
            max_attempts=2,
            initial_backoff_s=0.0,
            max_backoff_s=0.0,
            jitter_ratio=0.0,
            max_elapsed_time_s=None,
        ),
    )
    stub = _RetryQueryStub()
    client._stub = stub  # type: ignore[assignment]

    def call_metadata() -> Iterable[tuple[str, str]]:
        yield ("x-tenant", "acme")

    result = await client.query(doc_id=1, call_metadata=call_metadata())
    assert result.found is True
    assert stub.attempts == 2
    assert stub.call_metadata[0] == (("x-tenant", "acme"),)
    assert stub.call_metadata[1] == (("x-tenant", "acme"),)
    await client.close()


@pytest.mark.asyncio
async def test_async_circuit_breaker_fails_fast_after_threshold() -> None:
    client = AsyncKyroDBClient(
        target="127.0.0.1:50051",
        retry_policy=RetryPolicy(
            max_attempts=1,
            initial_backoff_s=0.0,
            max_backoff_s=0.0,
            jitter_ratio=0.0,
            max_elapsed_time_s=None,
        ),
        circuit_breaker_policy=CircuitBreakerPolicy(
            failure_threshold=1,
            recovery_timeout_s=60.0,
            half_open_success_threshold=1,
        ),
    )
    stub = _AlwaysUnavailableQueryStub()
    client._stub = stub  # type: ignore[assignment]

    with pytest.raises(ServiceUnavailableError):
        await client.query(doc_id=1)
    assert stub.attempts == 1

    with pytest.raises(CircuitOpenError):
        await client.query(doc_id=1)
    assert stub.attempts == 1
    await client.close()


@pytest.mark.asyncio
async def test_async_circuit_breaker_blocks_api_key_provider_side_effects_when_open() -> None:
    client = AsyncKyroDBClient(
        target="127.0.0.1:50051",
        retry_policy=RetryPolicy(
            max_attempts=1,
            initial_backoff_s=0.0,
            max_backoff_s=0.0,
            jitter_ratio=0.0,
            max_elapsed_time_s=None,
        ),
        circuit_breaker_policy=CircuitBreakerPolicy(
            failure_threshold=1,
            recovery_timeout_s=60.0,
            half_open_success_threshold=1,
        ),
    )
    stub = _AlwaysUnavailableQueryStub()
    client._stub = stub  # type: ignore[assignment]

    with pytest.raises(ServiceUnavailableError):
        await client.query(doc_id=1)
    assert stub.attempts == 1

    provider_calls = {"count": 0}

    def provider() -> str:
        provider_calls["count"] += 1
        if provider_calls["count"] == 1:
            return "rotating-key-v1"
        raise RuntimeError("provider should not run while circuit is open")

    client.set_api_key_provider(provider)

    calls_before_open_query = provider_calls["count"]
    with pytest.raises(CircuitOpenError):
        await client.query(doc_id=1)
    assert stub.attempts == 1
    assert provider_calls["count"] == calls_before_open_query
    await client.close()


@pytest.mark.asyncio
async def test_async_unary_batch_operations_auto_split_large_id_sets() -> None:
    client = AsyncKyroDBClient(target="127.0.0.1:50051", max_unary_batch_size=2)
    captured = _CapturedCall()
    client._stub = _FakeStub(captured)  # type: ignore[assignment]

    bulk_query = await client.bulk_query(doc_ids=[1, 2, 3, 4, 5], include_embeddings=False)
    assert bulk_query.total_requested == 5
    assert captured.bulk_query_calls == 3

    batch_delete = await client.batch_delete_ids(doc_ids=[1, 2, 3, 4, 5])
    assert batch_delete.deleted_count == 5
    assert captured.batch_delete_calls == 3
    await client.close()


@pytest.mark.skipif(np is None, reason="NumPy not installed")
@pytest.mark.asyncio
async def test_async_insert_rejects_numpy_non_finite_values() -> None:
    client = AsyncKyroDBClient(target="127.0.0.1:50051")
    captured = _CapturedCall()
    client._stub = _FakeStub(captured)  # type: ignore[assignment]

    with pytest.raises(ValueError, match="embedding contains non-finite values"):
        await client.insert(doc_id=1, embedding=np.array([0.1, np.nan], dtype=np.float32))
    with pytest.raises(ValueError, match="embedding contains non-finite values"):
        await client.insert(doc_id=1, embedding=np.array([0.1, np.inf], dtype=np.float32))
    await client.close()
