from __future__ import annotations

import math
from collections.abc import AsyncIterable, AsyncIterator, Iterable
from dataclasses import dataclass, field

import grpc
import pytest

from kyrodb import AsyncKyroDBClient, InsertItem, SearchQuery, exact
from kyrodb._generated import kyrodb_pb2 as pb2
from kyrodb.errors import CircuitOpenError, ServiceUnavailableError
from kyrodb.models import (
    BulkLoadResult,
    ConfigResult,
    DeleteResult,
    FlushResult,
    HealthResult,
    MetricsResult,
    SearchResponse,
    SnapshotResult,
    UpdateMetadataResult,
)
from kyrodb.retry import CircuitBreakerPolicy


class _FakeRpcError(grpc.RpcError):
    def __init__(self, code: grpc.StatusCode, details: str) -> None:
        super().__init__()
        self._code = code
        self._details = details

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

    def details(self) -> str:
        return self._details


@dataclass
class _Captured:
    bulk_insert_seen: int = 0
    bulk_load_seen: int = 0
    search_request: pb2.SearchRequest | None = None
    bulk_search_requests: list[pb2.SearchRequest] = field(default_factory=list)
    metadata: tuple[tuple[str, str], ...] | None = None
    timeout: float | None = None


class _AsyncRichStub:
    def __init__(self, captured: _Captured, *, fail_bulk_search: bool = False) -> None:
        self._captured = captured
        self._fail_bulk_search = fail_bulk_search

    async def Insert(
        self,
        request: pb2.InsertRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.InsertResponse:
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.InsertResponse(
            success=True,
            inserted_at=111,
            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:
        count = 0
        async for _request in requests:
            count += 1
        self._captured.bulk_insert_seen = count
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.InsertResponse(
            success=True,
            inserted_at=222,
            total_inserted=count,
            total_failed=0,
            tier=pb2.InsertResponse.HOT_TIER,
        )

    async def BulkLoadHnsw(
        self,
        requests: AsyncIterable[pb2.InsertRequest],
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.BulkLoadResponse:
        count = 0
        async for _request in requests:
            count += 1
        self._captured.bulk_load_seen = count
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.BulkLoadResponse(
            success=True,
            total_loaded=count,
            total_failed=0,
            load_duration_ms=1.2,
            avg_insert_rate=200.0,
            peak_memory_bytes=4096,
        )

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

    async def UpdateMetadata(
        self,
        request: pb2.UpdateMetadataRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.UpdateMetadataResponse:
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.UpdateMetadataResponse(success=True, existed=request.doc_id == 10)

    async def Query(
        self,
        request: pb2.QueryRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.QueryResponse:
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.QueryResponse(
            found=True,
            doc_id=request.doc_id,
            embedding=[0.2, 0.3] if request.include_embedding else [],
            metadata={"tenant": "beta"},
            served_from=pb2.QueryResponse.CACHE,
        )

    async def Search(
        self,
        request: pb2.SearchRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.SearchResponse:
        self._captured.search_request = request
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.SearchResponse(
            results=[pb2.SearchResult(doc_id=10, score=0.95, metadata={"tenant": "beta"})],
            total_found=1,
            search_latency_ms=0.55,
            search_path=pb2.SearchResponse.CACHE_HIT,
        )

    def BulkSearch(
        self,
        requests: AsyncIterable[pb2.SearchRequest],
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> AsyncIterator[pb2.SearchResponse]:
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())

        async def _stream() -> AsyncIterator[pb2.SearchResponse]:
            async for request in requests:
                self._captured.bulk_search_requests.append(request)
            if self._fail_bulk_search:
                raise _FakeRpcError(grpc.StatusCode.UNAVAILABLE, "x-api-key: async-secret")
            for index, _request in enumerate(self._captured.bulk_search_requests):
                yield pb2.SearchResponse(
                    results=[pb2.SearchResult(doc_id=index + 1, score=0.5)],
                    total_found=1,
                    search_latency_ms=0.25,
                    search_path=pb2.SearchResponse.COLD_TIER_ONLY,
                )

        return _stream()

    async def BulkQuery(
        self,
        request: pb2.BulkQueryRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.BulkQueryResponse:
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.BulkQueryResponse(
            results=[
                pb2.QueryResponse(
                    found=True,
                    doc_id=doc_id,
                    metadata={"batch": "true"},
                    served_from=pb2.QueryResponse.COLD_TIER,
                )
                for doc_id in request.doc_ids
            ],
            total_found=len(request.doc_ids),
            total_requested=len(request.doc_ids),
        )

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

    async def Health(
        self,
        request: pb2.HealthRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.HealthResponse:
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.HealthResponse(
            status=pb2.HealthResponse.HEALTHY,
            version="0.1.0",
            components={"component": request.component or "ok"},
            uptime_seconds=99,
            git_commit="def456",
        )

    async def Metrics(
        self,
        request: pb2.MetricsRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.MetricsResponse:
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.MetricsResponse(
            cache_hits=15,
            cache_misses=5,
            cache_hit_rate=0.75,
            cache_size=200,
            hot_tier_hits=12,
            hot_tier_misses=3,
            hot_tier_hit_rate=0.8,
            hot_tier_size=120,
            hot_tier_flushes=2,
            cold_tier_searches=7,
            cold_tier_size=2000,
            p50_latency_ms=1.1,
            p95_latency_ms=2.1,
            p99_latency_ms=3.1,
            total_queries=60,
            total_inserts=30,
            queries_per_second=120.0,
            inserts_per_second=55.0,
            memory_usage_bytes=3072,
            disk_usage_bytes=6144,
            cpu_usage_percent=22.0,
            overall_hit_rate=0.78,
            collected_at=5678,
        )

    async def FlushHotTier(
        self,
        request: pb2.FlushRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.FlushResponse:
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.FlushResponse(success=True, documents_flushed=4, flush_duration_ms=0.9)

    async def CreateSnapshot(
        self,
        request: pb2.SnapshotRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.SnapshotResponse:
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.SnapshotResponse(
            success=True,
            snapshot_path=request.path,
            documents_snapshotted=11,
            snapshot_size_bytes=16384,
        )

    async def GetConfig(
        self,
        request: pb2.ConfigRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.ConfigResponse:
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.ConfigResponse(
            hot_tier_max_size=100_000,
            hot_tier_max_age_seconds=600,
            hnsw_max_elements=1_000_000,
            data_dir="/tmp/kyrodb",
            fsync_policy="normal",
            snapshot_interval=1_000,
            flush_interval_seconds=15,
            embedding_dimension=4,
            version="0.1.0",
            hnsw_m=32,
            hnsw_ef_construction=400,
            hnsw_ef_search=128,
            hnsw_distance="cosine",
            hnsw_disable_normalization_check=False,
        )


class _AsyncInsertErrorStub:
    async def Insert(
        self,
        request: pb2.InsertRequest,
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.InsertResponse:
        raise _FakeRpcError(grpc.StatusCode.UNAVAILABLE, "x-api-key: direct-async-secret")


@pytest.mark.asyncio
async def test_async_context_manager_closes_channel(monkeypatch: pytest.MonkeyPatch) -> None:
    class _DummyChannel:
        def __init__(self) -> None:
            self.closed = False

        async def close(self) -> None:
            self.closed = True

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

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

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

    async with AsyncKyroDBClient(target="127.0.0.1:50051"):
        pass

    assert channel.closed is True


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

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

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

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

    client = AsyncKyroDBClient(target="127.0.0.1:50051")
    await client.wait_for_ready(timeout_s=None)
    await client.close()


@pytest.mark.asyncio
async def test_async_full_surface_returns_typed_models() -> None:
    client = AsyncKyroDBClient(target="127.0.0.1:50051", api_key="k_test", default_timeout_s=30.0)
    captured = _Captured()
    client._stub = _AsyncRichStub(captured)  # type: ignore[assignment]

    async def item_stream() -> AsyncIterator[InsertItem]:
        yield InsertItem.from_parts(doc_id=10, embedding=[0.1, 0.2], metadata={"a": "1"})
        yield InsertItem.from_parts(doc_id=11, embedding=[0.3, 0.4], metadata={"b": "2"})

    bulk_insert_ack = await client.bulk_insert(
        item_stream(),
        call_metadata=(("x-request-id", "async-1"),),
    )
    assert bulk_insert_ack.success is True
    assert captured.bulk_insert_seen == 2
    assert captured.metadata is not None
    assert ("x-request-id", "async-1") in captured.metadata

    bulk_load = await client.bulk_load_hnsw(
        [InsertItem.from_parts(doc_id=12, embedding=[0.5, 0.6])]
    )
    assert isinstance(bulk_load, BulkLoadResult)
    assert bulk_load.total_loaded == 1

    update = await client.update_metadata(doc_id=10, metadata={"tenant": "beta"}, merge=False)
    assert isinstance(update, UpdateMetadataResult)
    assert update.success is True

    query = await client.query(doc_id=10, include_embedding=True, namespace="default")
    assert query.found is True
    assert query.embedding is not None
    assert query.embedding == pytest.approx((0.2, 0.3))

    search = await client.search(
        query_embedding=[0.1, 0.2],
        k=5,
        namespace="default",
        filter=exact("tenant", "beta"),
        ef_search=16,
    )
    assert isinstance(search, SearchResponse)
    assert search.results[0].doc_id == 10
    assert captured.search_request is not None
    assert captured.search_request.filter.exact.key == "tenant"

    async def request_stream() -> AsyncIterator[SearchQuery]:
        yield SearchQuery.from_parts(query_embedding=[0.1, 0.2], k=1, namespace="default")
        yield SearchQuery.from_parts(query_embedding=[0.3, 0.4], k=1, namespace="default")

    bulk_search_results = [result async for result in client.bulk_search(request_stream())]
    assert len(bulk_search_results) == 2
    assert len(captured.bulk_search_requests) == 2

    health = await client.health(component="storage")
    assert isinstance(health, HealthResult)
    assert health.status == "HEALTHY"

    metrics = await client.metrics(categories=["cache"])
    assert isinstance(metrics, MetricsResult)
    assert metrics.cache_hits == 15

    flush = await client.flush_hot_tier(force=False)
    assert isinstance(flush, FlushResult)
    assert flush.documents_flushed == 4

    snapshot = await client.create_snapshot(path="/tmp/async-snapshot")
    assert isinstance(snapshot, SnapshotResult)
    assert snapshot.snapshot_path == "/tmp/async-snapshot"

    config = await client.get_config()
    assert isinstance(config, ConfigResult)
    assert config.hnsw_m == 32

    await client.close()


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

    result = await client.delete(doc_id=10, namespace="default")
    assert isinstance(result, DeleteResult)
    assert result.success is True
    assert result.existed is True

    await client.close()


@pytest.mark.asyncio
async def test_async_surface_validates_inputs() -> None:
    client = AsyncKyroDBClient(target="127.0.0.1:50051")
    client._stub = _AsyncRichStub(_Captured())  # type: ignore[assignment]

    with pytest.raises(ValueError, match="doc_id exceeds uint64 range"):
        await client.insert(doc_id=0xFFFFFFFFFFFFFFFF + 1, embedding=[0.1])
    with pytest.raises(ValueError, match="must be non-empty"):
        await client.insert(doc_id=1, embedding=[])
    with pytest.raises(ValueError, match="non-finite"):
        await client.insert(doc_id=1, embedding=[math.nan])
    with pytest.raises(ValueError, match="k must be between 1 and 1000"):
        await client.search(query_embedding=[0.1], k=0)
    with pytest.raises(ValueError, match="ef_search must be >= 0"):
        await client.search(query_embedding=[0.1], k=1, ef_search=-1)
    with pytest.raises(ValueError, match="doc_ids must be non-empty"):
        await client.bulk_query(doc_ids=[])
    with pytest.raises(ValueError, match="doc_ids must be non-empty"):
        await client.batch_delete_ids(doc_ids=[])

    await client.close()


@pytest.mark.asyncio
async def test_async_bulk_search_maps_rpc_error() -> None:
    client = AsyncKyroDBClient(target="127.0.0.1:50051")
    client._stub = _AsyncRichStub(
        _Captured(),
        fail_bulk_search=True,
    )  # type: ignore[assignment]

    async def request_stream() -> AsyncIterator[SearchQuery]:
        yield SearchQuery.from_parts(query_embedding=[0.1], k=1)

    with pytest.raises(ServiceUnavailableError):
        _ = [result async for result in client.bulk_search(request_stream())]

    await client.close()


@pytest.mark.asyncio
async def test_async_bulk_search_respects_circuit_breaker() -> None:
    captured = _Captured()
    client = AsyncKyroDBClient(
        target="127.0.0.1:50051",
        circuit_breaker_policy=CircuitBreakerPolicy(
            failure_threshold=1,
            recovery_timeout_s=60.0,
            half_open_success_threshold=1,
        ),
    )
    client._stub = _AsyncRichStub(captured, fail_bulk_search=True)  # type: ignore[assignment]

    async def request_stream() -> AsyncIterator[SearchQuery]:
        yield SearchQuery.from_parts(query_embedding=[0.1], k=1)

    with pytest.raises(ServiceUnavailableError):
        _ = [result async for result in client.bulk_search(request_stream())]
    first_failure_count = len(captured.bulk_search_requests)

    with pytest.raises(CircuitOpenError):
        _ = [result async for result in client.bulk_search(request_stream())]
    assert len(captured.bulk_search_requests) == first_failure_count
    await client.close()


@pytest.mark.asyncio
async def test_async_search_logs_dimension_mismatch(caplog: pytest.LogCaptureFixture) -> None:
    client = AsyncKyroDBClient(target="127.0.0.1:50051")
    captured = _Captured()
    client._stub = _AsyncRichStub(captured)  # type: ignore[assignment]

    caplog.set_level("DEBUG", logger="kyrodb")
    await client.insert(doc_id=1, embedding=[0.1, 0.2], namespace="default")
    await client.search(query_embedding=[0.1, 0.2, 0.3], k=1, namespace="default")

    assert "Search embedding dimension differs from last insert" in caplog.text
    await client.close()


@pytest.mark.asyncio
async def test_async_safe_call_handles_broken_key_provider_when_logging_error() -> None:
    state = {"calls": 0}

    def provider() -> str:
        state["calls"] += 1
        if state["calls"] >= 3:
            raise RuntimeError("provider failure")
        return "key_ok"

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

    with pytest.raises(ServiceUnavailableError):
        await client.insert(doc_id=1, embedding=[0.1, 0.2])

    await client.close()
