from __future__ import annotations

import math
from collections.abc import Iterable, Iterator
from dataclasses import dataclass, field

import grpc
import pytest

from kyrodb import InsertItem, KyroDBClient, SearchQuery, exact
from kyrodb._generated import kyrodb_pb2 as pb2
from kyrodb.errors import CircuitOpenError, ServiceUnavailableError
from kyrodb.models import (
    BulkLoadResult,
    ConfigResult,
    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 _RichStub:
    def __init__(self, captured: _Captured, *, fail_bulk_search: bool = False) -> None:
        self._captured = captured
        self._fail_bulk_search = fail_bulk_search

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

    def BulkInsert(
        self,
        requests: Iterable[pb2.InsertRequest],
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.InsertResponse:
        self._captured.bulk_insert_seen = len(list(requests))
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.InsertResponse(
            success=True,
            inserted_at=222,
            total_inserted=self._captured.bulk_insert_seen,
            total_failed=0,
            tier=pb2.InsertResponse.HOT_TIER,
        )

    def BulkLoadHnsw(
        self,
        requests: Iterable[pb2.InsertRequest],
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> pb2.BulkLoadResponse:
        self._captured.bulk_load_seen = len(list(requests))
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        return pb2.BulkLoadResponse(
            success=True,
            total_loaded=self._captured.bulk_load_seen,
            total_failed=0,
            load_duration_ms=1.5,
            avg_insert_rate=300.0,
            peak_memory_bytes=2048,
        )

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

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

    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.1, 0.2] if request.include_embedding else [],
            metadata={"tenant": "acme"},
            served_from=pb2.QueryResponse.HOT_TIER,
        )

    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=7, score=0.99, metadata={"tenant": "acme"})],
            total_found=1,
            search_latency_ms=0.4,
            search_path=pb2.SearchResponse.HOT_TIER_ONLY,
        )

    def BulkSearch(
        self,
        requests: Iterable[pb2.SearchRequest],
        timeout: float | None = None,
        metadata: Iterable[tuple[str, str]] | None = None,
    ) -> Iterator[pb2.SearchResponse]:
        self._captured.bulk_search_requests = list(requests)
        self._captured.timeout = timeout
        self._captured.metadata = tuple(metadata or ())
        if self._fail_bulk_search:
            raise _FakeRpcError(grpc.StatusCode.UNAVAILABLE, "x-api-key: bulk-search-secret")
        return iter(
            [
                pb2.SearchResponse(
                    results=[pb2.SearchResult(doc_id=1, score=0.8)],
                    total_found=1,
                    search_latency_ms=0.3,
                    search_path=pb2.SearchResponse.CACHE_HIT,
                ),
                pb2.SearchResponse(
                    results=[pb2.SearchResult(doc_id=2, score=0.7)],
                    total_found=1,
                    search_latency_ms=0.35,
                    search_path=pb2.SearchResponse.HOT_AND_COLD,
                ),
            ]
        )

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

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

    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={"storage": request.component or "ok"},
            uptime_seconds=42,
            git_commit="abc123",
        )

    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=10,
            cache_misses=2,
            cache_hit_rate=0.83,
            cache_size=100,
            hot_tier_hits=9,
            hot_tier_misses=1,
            hot_tier_hit_rate=0.9,
            hot_tier_size=80,
            hot_tier_flushes=1,
            cold_tier_searches=5,
            cold_tier_size=1000,
            p50_latency_ms=1.0,
            p95_latency_ms=2.0,
            p99_latency_ms=3.0,
            total_queries=50,
            total_inserts=20,
            queries_per_second=100.0,
            inserts_per_second=40.0,
            memory_usage_bytes=2048,
            disk_usage_bytes=4096,
            cpu_usage_percent=12.5,
            overall_hit_rate=0.84,
            collected_at=1234,
        )

    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=3, flush_duration_ms=0.7)

    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=8,
            snapshot_size_bytes=8192,
        )

    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 _InsertErrorStub:
    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-secret")


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

    bulk_insert_ack = 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"}),
        ],
        call_metadata=(("x-request-id", "sync-1"),),
    )
    assert bulk_insert_ack.success is True
    assert captured.bulk_insert_seen == 2
    assert captured.metadata is not None
    assert ("x-request-id", "sync-1") in captured.metadata

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

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

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

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

    bulk_search = list(
        client.bulk_search(
            [SearchQuery.from_parts(query_embedding=[0.1, 0.2], k=1, namespace="default")],
            timeout_s=5.0,
        )
    )
    assert len(bulk_search) == 2
    assert captured.bulk_search_requests[0].k == 1

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

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

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

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

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

    client.close()


def test_sync_surface_validates_inputs() -> None:
    client = KyroDBClient(target="127.0.0.1:50051")
    client._stub = _RichStub(_Captured())  # type: ignore[assignment]

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

    client.close()


def test_sync_bulk_search_maps_rpc_error() -> None:
    client = KyroDBClient(target="127.0.0.1:50051")
    client._stub = _RichStub(_Captured(), fail_bulk_search=True)  # type: ignore[assignment]

    with pytest.raises(ServiceUnavailableError):
        list(client.bulk_search([SearchQuery.from_parts(query_embedding=[0.1], k=1)]))

    client.close()


def test_sync_bulk_search_respects_circuit_breaker() -> None:
    captured = _Captured()
    client = KyroDBClient(
        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 = _RichStub(captured, fail_bulk_search=True)  # type: ignore[assignment]

    with pytest.raises(ServiceUnavailableError):
        list(client.bulk_search([SearchQuery.from_parts(query_embedding=[0.1], k=1)]))
    first_failure_count = len(captured.bulk_search_requests)

    with pytest.raises(CircuitOpenError):
        list(client.bulk_search([SearchQuery.from_parts(query_embedding=[0.1], k=1)]))
    assert len(captured.bulk_search_requests) == first_failure_count
    client.close()


def test_sync_search_logs_dimension_mismatch(caplog: pytest.LogCaptureFixture) -> None:
    client = KyroDBClient(target="127.0.0.1:50051")
    captured = _Captured()
    client._stub = _RichStub(captured)  # type: ignore[assignment]

    caplog.set_level("DEBUG", logger="kyrodb")
    client.insert(doc_id=1, embedding=[0.1, 0.2], namespace="default")
    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
    client.close()


def test_sync_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 = KyroDBClient(target="127.0.0.1:50051")
    client._stub = _InsertErrorStub()  # type: ignore[assignment]
    client.set_api_key_provider(provider)

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

    client.close()
