# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""
Implementation of vLLM KV cache manager protocol.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Optional

import torch
from typing_extensions import override
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
    KVConnectorBase_V1,
    KVConnectorMetadata,
    KVConnectorRole,
)
from vllm.v1.core.sched.output import SchedulerOutput

if TYPE_CHECKING:
    from vllm.attention.backends.abstract import AttentionMetadata
    from vllm.config import VllmConfig
    from vllm.forward_context import ForwardContext
    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
    from vllm.v1.kv_cache_interface import KVCacheConfig
    from vllm.v1.request import Request


# from kvbm.vllm_integration.kv_cache_utils import KvbmCacheBlocks
from kvbm.utils import nvtx_annotate
from kvbm.vllm_integration.connector_leader import KvConnectorLeader
from kvbm.vllm_integration.connector_worker import KvConnectorWorker

EngineId = str


class DynamoConnectorMetadata(KVConnectorMetadata):
    def __init__(self, metadata: bytes):
        assert isinstance(metadata, bytes)
        self.metadata = metadata


class DynamoConnector(KVConnectorBase_V1):
    def __init__(
        self,
        vllm_config: "VllmConfig",
        role: KVConnectorRole,
        kv_cache_config: Optional["KVCacheConfig"] = None,
    ):
        super().__init__(
            vllm_config=vllm_config, role=role, kv_cache_config=kv_cache_config
        )

        assert vllm_config.kv_transfer_config is not None
        assert vllm_config.kv_transfer_config.engine_id is not None
        self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id

        if role == KVConnectorRole.SCHEDULER:
            self._scheduler = KvConnectorLeader(
                vllm_config=vllm_config, engine_id=self.engine_id
            )
            self._worker = None
        elif role == KVConnectorRole.WORKER:
            self._worker = KvConnectorWorker(
                vllm_config=vllm_config, engine_id=self.engine_id
            )
            self._scheduler = None

    # Scheduler/Leader

    @nvtx_annotate(category="scheduler")
    def get_num_new_matched_tokens(
        self,
        request: "Request",
        num_computed_tokens: int,
    ) -> tuple[int, bool]:
        return self._scheduler.get_num_new_matched_tokens(request, num_computed_tokens)

    @nvtx_annotate(category="scheduler")
    def update_state_after_alloc(
        self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
    ):
        self._scheduler.update_state_after_alloc(request, blocks, num_external_tokens)

    @nvtx_annotate(category="scheduler")
    def build_connector_meta(
        self, scheduler_output: SchedulerOutput
    ) -> KVConnectorMetadata:
        data = self._scheduler.build_connector_meta(scheduler_output)
        return DynamoConnectorMetadata(data)

    @nvtx_annotate(category="scheduler")
    def request_finished(
        self,
        request: "Request",
        block_ids: list[int],
    ) -> tuple[bool, Optional[dict[str, Any]]]:
        return self._scheduler.request_finished(request, block_ids)

    # Worker

    @nvtx_annotate(category="worker")
    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
        self._worker.register_kv_caches(kv_caches)

    @nvtx_annotate(category="worker")
    @override
    def bind_connector_metadata(
        self, connector_metadata: DynamoConnectorMetadata
    ) -> None:
        # Must call super() to set _connector_metadata so has_connector_metadata() returns True
        # This is required for save_kv_layer to be called during the forward pass
        super().bind_connector_metadata(connector_metadata)
        assert isinstance(connector_metadata.metadata, bytes)
        self._worker.bind_connector_metadata(connector_metadata.metadata)

    @nvtx_annotate(category="worker")
    @override
    def clear_connector_metadata(self) -> None:
        super().clear_connector_metadata()
        self._worker.clear_connector_metadata()

    @nvtx_annotate(category="worker")
    @override
    def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
        self._worker.start_load_kv(forward_context, **kwargs)

    @nvtx_annotate(category="worker")
    @override
    def wait_for_layer_load(self, layer_name: str) -> None:
        pass

    @nvtx_annotate(category="worker")
    @override
    def save_kv_layer(
        self,
        layer_name: str,
        kv_layer: torch.Tensor,
        attn_metadata: "AttentionMetadata",
        **kwargs,
    ) -> None:
        self._worker.save_kv_layer(layer_name, kv_layer, attn_metadata, **kwargs)

    @nvtx_annotate(category="worker")
    @override
    def wait_for_save(self):
        pass

    def get_finished(
        self, finished_req_ids: set[str]
    ) -> tuple[Optional[set[str]], Optional[set[str]]]:
        return self._worker.get_finished(finished_req_ids)
