from __future__ import annotations

from collections.abc import Iterable
from dataclasses import dataclass

import grpc
import pytest

from kyrodb import KyroDBClient
from kyrodb._generated import kyrodb_pb2 as pb2
from kyrodb.errors import CircuitOpenError, ServiceUnavailableError
from kyrodb.models import BatchDeleteResult, BulkQueryResult, DeleteResult, InsertAck
from kyrodb.retry import CircuitBreakerPolicy, RetryPolicy

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


@dataclass
class _CapturedCall:
    insert_request: 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

    def Insert(
        self,
        request: pb2.InsertRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.InsertResponse:
        self._captured.insert_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,
        )

    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=True)

    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(
            results=[
                pb2.QueryResponse(
                    found=True,
                    doc_id=10,
                    metadata={"tenant": "acme"},
                    served_from=pb2.QueryResponse.HOT_TIER,
                )
            ],
            total_found=1,
            total_requested=len(request.doc_ids),
        )

    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], ...]] = []

    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

    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)


def test_sync_insert_returns_typed_ack_and_uses_default_timeout() -> None:
    client = KyroDBClient(target="127.0.0.1:50051", api_key="k_test", default_timeout_s=30.0)
    captured = _CapturedCall()
    client._stub = _FakeStub(captured)  # type: ignore[assignment]

    ack = client.insert(doc_id=1, embedding=[0.1, 0.2], metadata={"source": "unit-test"})

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

    client.close()


def test_sync_delete_allows_explicit_unbounded_timeout() -> None:
    client = KyroDBClient(target="127.0.0.1:50051", default_timeout_s=30.0)
    captured = _CapturedCall()
    client._stub = _FakeStub(captured)  # type: ignore[assignment]

    delete = client.delete(doc_id=7, timeout_s=None)

    assert isinstance(delete, DeleteResult)
    assert delete.success is True
    assert delete.existed is True
    assert captured.timeout is None

    client.close()


def test_sync_bulk_query_and_batch_delete_return_typed_models() -> None:
    client = KyroDBClient(target="127.0.0.1:50051")
    captured = _CapturedCall()
    client._stub = _FakeStub(captured)  # type: ignore[assignment]

    bulk_query = client.bulk_query(doc_ids=[10], include_embeddings=False, namespace="default")
    assert isinstance(bulk_query, BulkQueryResult)
    assert bulk_query.total_found == 1
    assert bulk_query.results[0].doc_id == 10
    assert bulk_query.results[0].served_from == "HOT_TIER"

    batch_delete = client.batch_delete_ids(doc_ids=[10, 11], namespace="default")
    assert isinstance(batch_delete, BatchDeleteResult)
    assert batch_delete.success is True
    assert batch_delete.deleted_count == 2

    client.close()


def test_sync_api_key_rotation_with_set_api_key() -> None:
    client = KyroDBClient(target="127.0.0.1:50051", api_key="key_v1")
    captured = _CapturedCall()
    client._stub = _FakeStub(captured)  # type: ignore[assignment]

    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")
    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

    client.close()


def test_sync_api_key_provider_is_used_per_call() -> None:
    key_state = {"value": "key_v1"}

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

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

    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"
    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

    client.close()


def test_sync_call_timeout_validates_positive_values() -> None:
    client = KyroDBClient(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"):
        client.insert(doc_id=1, embedding=[0.1, 0.2], timeout_s=0)
    client.close()


def test_sync_api_key_provider_requires_callable() -> None:
    client = KyroDBClient(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]
    client.close()


def test_sync_retry_preserves_call_metadata_for_generators() -> None:
    client = KyroDBClient(
        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 = 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"),)
    client.close()


def test_sync_circuit_breaker_fails_fast_after_threshold() -> None:
    client = KyroDBClient(
        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):
        client.query(doc_id=1)
    assert stub.attempts == 1

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


def test_sync_circuit_breaker_blocks_api_key_provider_side_effects_when_open() -> None:
    client = KyroDBClient(
        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):
        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):
        client.query(doc_id=1)
    assert stub.attempts == 1
    assert provider_calls["count"] == calls_before_open_query
    client.close()


def test_sync_unary_batch_operations_auto_split_large_id_sets() -> None:
    client = KyroDBClient(target="127.0.0.1:50051", max_unary_batch_size=2)
    captured = _CapturedCall()
    client._stub = _FakeStub(captured)  # type: ignore[assignment]

    bulk_query = 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 = client.batch_delete_ids(doc_ids=[1, 2, 3, 4, 5])
    assert batch_delete.deleted_count == 5
    assert captured.batch_delete_calls == 3
    client.close()


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

    ack = client.insert(doc_id=1, embedding=np.array([0.1, 0.2], dtype=np.float32))
    assert ack.success is True
    assert captured.insert_request is not None
    assert list(captured.insert_request.embedding) == pytest.approx([0.1, 0.2])
    client.close()


@pytest.mark.skipif(np is None, reason="NumPy not installed")
def test_sync_insert_rejects_numpy_non_finite_values() -> None:
    client = KyroDBClient(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"):
        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"):
        client.insert(doc_id=1, embedding=np.array([0.1, np.inf], dtype=np.float32))
    client.close()
