# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import copy
import queue
import threading
import time
from collections import defaultdict
from collections.abc import Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Optional

import torch
from vllm.config import VllmConfig, get_layers_from_vllm_config
from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole
from vllm.distributed.parallel_state import get_pp_group, get_tp_group
from vllm.logger import logger
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
from vllm.model_executor.layers.mamba.abstract import MambaBase
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec

from vllm_ascend.distributed.kv_transfer.kv_pool.cpu_offload.metadata import (
    MetadataServer,
    MetadataServerProc,
    MLAConfig,
)

if TYPE_CHECKING:
    from vllm.forward_context import ForwardContext
    from vllm.v1.attention.backend import AttentionMetadata  # type: ignore
    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
    from vllm.v1.kv_cache_interface import KVCacheConfig
    from vllm.v1.request import Request

from vllm.model_executor.layers.attention import Attention, MLAAttention


@dataclass
class ReqMeta:
    gpu_block_ids: list[int]
    cpu_block_ids: list[int]
    num_scheduled_tokens: int
    num_computed_tokens: int
    num_gpu_computed_tokens: int
    num_cpu_computed_tokens: int

    def update(self, other: "ReqMeta"):
        self.gpu_block_ids.extend(other.gpu_block_ids)
        self.cpu_block_ids.extend(other.cpu_block_ids)
        self.num_scheduled_tokens = other.num_scheduled_tokens
        self.num_computed_tokens = other.num_computed_tokens
        self.num_gpu_computed_tokens = other.num_gpu_computed_tokens
        self.num_cpu_computed_tokens = other.num_cpu_computed_tokens


@dataclass
class CPUOffloadingConnectorMetadata(KVConnectorMetadata):
    requests: dict[str, ReqMeta]
    finished_req_ids: set[str]


class CPUOffloadingConnector(KVConnectorBase_V1):
    def __init__(
        self, vllm_config: VllmConfig, role: KVConnectorRole, kv_cache_config: Optional["KVCacheConfig"] = None
    ):
        self._connector_metadata = CPUOffloadingConnectorMetadata(requests={}, finished_req_ids=set())
        if not vllm_config.cache_config.enable_prefix_caching:
            self.connector_scheduler: CPUOffloadingConnectorScheduler | None = None
            self.connector_worker: CPUOffloadingConnectorWorker | None = None
        elif role == KVConnectorRole.SCHEDULER:
            self.connector_scheduler = CPUOffloadingConnectorScheduler(vllm_config)
            self.connector_worker = None
        elif role == KVConnectorRole.WORKER:
            self.connector_scheduler = None
            self.connector_worker = CPUOffloadingConnectorWorker(vllm_config)

    # ==============================
    # Worker-side methods
    # ==============================

    def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> None:
        if self.connector_worker is not None:
            assert isinstance(connector_metadata, CPUOffloadingConnectorMetadata)
            self.connector_worker.bind_connector_metadata(connector_metadata)

    def clear_connector_metadata(self) -> None:
        assert self.connector_worker is not None
        self.connector_worker.clear_connector_metadata()

    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
        if self.connector_worker is not None:
            self.connector_worker.register_kv_caches(kv_caches)

    def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
        if self.connector_worker is not None:
            self.connector_worker.start_load_kv()

    def wait_for_layer_load(self, layer_name: str) -> None:
        if self.connector_worker is not None:
            self.connector_worker.wait_for_layer_load()

    def save_kv_layer(
        self, layer_name: str, kv_layer: torch.Tensor, attn_metadata: "AttentionMetadata", **kwargs
    ) -> None:
        pass

    def wait_for_save(self):
        pass

    def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str] | None, set[str] | None]:
        assert self.connector_worker is not None
        return self.connector_worker.get_finished(), None

    # Scheduler-side methods
    # ==============================

    def get_num_new_matched_tokens(self, request: "Request", num_computed_tokens: int) -> tuple[int, bool]:
        if self.connector_scheduler is not None:
            return self.connector_scheduler.get_num_new_matched_tokens(request, num_computed_tokens)
        return 0, False

    def update_state_after_alloc(self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int):
        if self.connector_scheduler is not None:
            return self.connector_scheduler.update_state_after_alloc(request)

    def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata:
        if self.connector_scheduler is not None:
            return self.connector_scheduler.build_connector_meta(scheduler_output)
        return KVConnectorMetadata()

    def request_finished(self, request: "Request", block_ids: list[int]) -> tuple[bool, dict[str, Any] | None]:
        if self.connector_scheduler is not None:
            self.connector_scheduler.request_finished(request)
        return True, None


class CPUOffloadingConnectorScheduler:
    def __init__(self, vllm_config: VllmConfig):
        logger.info("init CPUOffloadingConnectorScheduler")
        self.vllm_config = vllm_config
        self.block_size = vllm_config.cache_config.block_size
        self.use_mla = vllm_config.model_config.use_mla
        self.num_gpu_computed_tokens: dict[str, int] = {}
        self.num_cpu_computed_tokens: dict[str, int] = {}
        self.allocated_req_ids: set[str] = set()
        self.finished_req_ids: list[str] = []
        self.zmq_rpc_client = MetadataServer.ZMQRPCClient()
        self.zmq_rpc_client.call("post_init")
        if vllm_config.kv_transfer_config is not None:
            self.swap_in_threshold = vllm_config.kv_transfer_config.get_from_extra_config("swap_in_threshold", 0)
        else:
            self.swap_in_threshold = 0
        logger.info(f"swap_in_threshold: {self.swap_in_threshold}")

    def get_num_new_matched_tokens(self, ori_request: "Request", num_computed_tokens: int) -> tuple[int, bool]:
        request = copy.deepcopy(ori_request)
        request.get_hash_new_full_blocks = None
        num_cpu_computed_tokens, load_async = self.zmq_rpc_client.call("get_matched_num_and_touch", request)
        self.num_gpu_computed_tokens[request.request_id] = num_computed_tokens
        self.num_cpu_computed_tokens[request.request_id] = num_cpu_computed_tokens
        if num_cpu_computed_tokens - num_computed_tokens >= self.swap_in_threshold:
            return num_cpu_computed_tokens - num_computed_tokens, load_async
        else:
            return 0, load_async

    def update_state_after_alloc(self, request: "Request"):
        self.allocated_req_ids.add(request.request_id)

    def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata:
        num_tokens = {}
        # process scheduled_new_reqs
        for req in scheduler_output.scheduled_new_reqs:
            req_id = req.req_id
            num_tokens[req_id] = req.num_computed_tokens + scheduler_output.num_scheduled_tokens[req_id]

        # process scheduled_cached_reqs
        cached_reqs = scheduler_output.scheduled_cached_reqs
        for idx, req_id in enumerate(cached_reqs.req_ids):
            num_tokens[req_id] = cached_reqs.num_computed_tokens[idx] + scheduler_output.num_scheduled_tokens[req_id]

        unallocated_req_ids = set(
            self.num_gpu_computed_tokens.keys() - self.allocated_req_ids - scheduler_output.num_scheduled_tokens.keys()
        )
        new_cpu_block_ids = self.zmq_rpc_client.call("allocate_slots", num_tokens, unallocated_req_ids)
        metadata = CPUOffloadingConnectorMetadata(
            requests={},
            finished_req_ids=set(self.finished_req_ids),
        )
        for req in scheduler_output.scheduled_new_reqs:
            req_id = req.req_id
            gpu_block_ids = req.block_ids[0]
            metadata.requests[req_id] = ReqMeta(
                gpu_block_ids=[] if gpu_block_ids is None else gpu_block_ids,
                cpu_block_ids=new_cpu_block_ids.get(req_id, []),
                num_scheduled_tokens=scheduler_output.num_scheduled_tokens[req_id],
                num_computed_tokens=req.num_computed_tokens,
                num_gpu_computed_tokens=self.num_gpu_computed_tokens[req_id],
                num_cpu_computed_tokens=self.num_cpu_computed_tokens[req_id],
            )

        for idx, req_id in enumerate(cached_reqs.req_ids):
            gpu_block_ids = cached_reqs.new_block_ids[idx]
            metadata.requests[req_id] = ReqMeta(
                gpu_block_ids=[] if gpu_block_ids is None else gpu_block_ids,
                cpu_block_ids=new_cpu_block_ids.get(req_id, []),
                num_scheduled_tokens=scheduler_output.num_scheduled_tokens[req_id],
                num_computed_tokens=cached_reqs.num_computed_tokens[idx],
                num_gpu_computed_tokens=cached_reqs.num_computed_tokens[idx],
                num_cpu_computed_tokens=cached_reqs.num_computed_tokens[idx],
            )
        self.num_gpu_computed_tokens.clear()
        self.num_cpu_computed_tokens.clear()
        self.allocated_req_ids.clear()
        self.finished_req_ids.clear()
        return metadata

    def request_finished(self, ori_request: "Request"):
        request = copy.deepcopy(ori_request)
        request.get_hash_new_full_blocks = None
        self.finished_req_ids.append(request.request_id)
        # inform metadata server to record request, and free it after finish sending
        self.zmq_rpc_client.call("record_request_cache_and_free_slots", request)


class CPUOffloadingConnectorWorker:
    def __init__(self, vllm_config: VllmConfig):
        logger.info("init CPUOffloadingConnectorWorker")
        self.vllm_config = vllm_config
        self.block_size = vllm_config.cache_config.block_size
        self.pp_rank = get_pp_group().rank_in_group
        self.tp_group = get_tp_group()
        self.tp_rank = self.tp_group.rank_in_group
        self.tp_world_size = self.tp_group.world_size
        self.use_mla = vllm_config.model_config.use_mla

        self.requests: dict[str, ReqMeta] = {}
        self.load_stream = torch.npu.Stream()
        self.save_stream = torch.npu.Stream()
        self.zmq_rpc_client = MetadataServer.ZMQRPCClient()
        self.load_block_mapping: list[tuple[int, int]] = []
        self.save_input_queue: queue.Queue[tuple[str, ReqMeta]] = queue.Queue()
        self.save_output_queue: queue.Queue[str] = queue.Queue()
        self.save_thread = threading.Thread(target=self._save_listener)
        self.save_thread.start()
        self.done_sending_count: defaultdict[str, int] = defaultdict(int)

        # start metadata server to init cpu_kv_cache_manager and handle rpc requests
        # all dp shared the same metadata server, only start the process on data_rank 0
        if vllm_config.parallel_config.data_parallel_rank == 0 and self.tp_rank == 0 and self.pp_rank == 0:
            config = VllmConfig()
            config.cache_config = vllm_config.cache_config
            config.parallel_config = vllm_config.parallel_config
            config.kv_transfer_config = vllm_config.kv_transfer_config
            self.init_metadata_server(config)
        self._wait_for_metadata_process_start()

    def init_metadata_server(self, vllm_config: VllmConfig):
        self.metadata_thread = threading.Thread(
            target=MetadataServerProc.run_metadata_server,
            args=(vllm_config,),
        )
        self.metadata_thread.daemon = True
        self.metadata_thread.start()

    def _wait_for_metadata_process_start(self):
        # TODO: wait for metadata server to start, add a rpc to check if ready
        while True:
            try:
                if self.zmq_rpc_client.call("ready"):
                    break
            except Exception as e:
                logger.info(f"wait for metadata server to start, error: {e}")
                time.sleep(1)

    def bind_connector_metadata(self, connector_metadata: CPUOffloadingConnectorMetadata) -> None:
        for req_id, req in connector_metadata.requests.items():
            if req_id in self.requests:
                self.requests[req_id].update(req)
                req = self.requests[req_id]
            else:
                self.requests[req_id] = req
            for i in range(req.num_gpu_computed_tokens // self.block_size, req.num_computed_tokens // self.block_size):
                self.load_block_mapping.append((req.cpu_block_ids[i], req.gpu_block_ids[i]))
        for req_id in connector_metadata.finished_req_ids:
            if req_id in self.requests:
                self.save_input_queue.put((req_id, self.requests[req_id]))

    def clear_connector_metadata(self) -> None:
        self.load_block_mapping.clear()

    def register_kv_caches(self, kv_caches: dict[str, Sequence[torch.Tensor]]):
        self.gpu_kv_caches = kv_caches
        model_config = self.vllm_config.model_config
        mla_config: MLAConfig | None = None
        if model_config.use_mla:
            mla_config = MLAConfig(
                model_config.hf_text_config.kv_lora_rank, model_config.hf_text_config.qk_rope_head_dim
            )
        self.cpu_kv_caches = list(
            self.zmq_rpc_client.call(
                "init_cpu_kv_caches",
                self.pp_rank,
                self.tp_rank,
                get_kv_cache_spec(self.vllm_config),
                mla_config,
            ).values()
        )

    def start_load_kv(self) -> None:
        self.current_layer = 0
        self.gpu_kv_caches_load_iter = iter(self.gpu_kv_caches.values())
        self.load_kv_layer(0)

    def wait_for_layer_load(self) -> None:
        # TODO: Replace with `torch.npu.current_stream().wait_stream(self.load_stream)` after fixing the bug.
        self.load_stream.synchronize()
        self.current_layer += 1
        self.load_kv_layer(self.current_layer)

    def load_kv_layer(self, layer: int):
        if layer == len(self.gpu_kv_caches):
            return
        gpu_kv_caches = next(self.gpu_kv_caches_load_iter)
        cpu_kv_caches = self.cpu_kv_caches[layer]
        with torch.npu.stream(self.load_stream):
            for cpu_block_id, gpu_block_id in self.load_block_mapping:
                for gpu_layer_part, cpu_layer_part in zip(gpu_kv_caches, cpu_kv_caches):
                    gpu_layer_part[gpu_block_id].copy_(cpu_layer_part[cpu_block_id], non_blocking=True)

    def get_finished(self) -> set[str]:
        done_sending: set[str] = set()
        while True:
            try:
                id = self.save_output_queue.get_nowait()
            except queue.Empty:
                break
            done_sending.add(id)
        for id in done_sending:
            del self.requests[id]
        if self.tp_world_size == 1:
            return done_sending
        if self.tp_rank == 0:
            for req_id in done_sending:
                self.done_sending_count[req_id] += 1
            other_ranks_finished_ids: list[str] = []
            for i in range(1, self.tp_world_size):
                other_ranks_finished_ids.extend(self.tp_group.recv_object(src=i))
            for req_id in other_ranks_finished_ids:
                self.done_sending_count[req_id] += 1
            all_done_sending: set[str] = set()
            for req_id in list(self.done_sending_count.keys()):
                if self.done_sending_count[req_id] == self.tp_world_size:
                    del self.done_sending_count[req_id]
                    all_done_sending.add(req_id)
            # release cpu_kv_cache after request sending finished
            # to avoid rpc blocking, use thread to call rpc asynchronously
            sending_finished_thread = threading.Thread(target=self._sending_finished, args=(all_done_sending,))
            sending_finished_thread.daemon = True
            sending_finished_thread.start()

            return all_done_sending
        else:
            self.tp_group.send_object(done_sending, dst=0)
            return done_sending

    def _sending_finished(self, all_done_sending):
        for req_id in all_done_sending:
            logger.debug(f"call cache_and_free_slots for req_id: {req_id}")
            self.zmq_rpc_client.call("cache_and_free_slots", req_id)

    def _save_listener(self):
        save_block_mapping = []
        while True:
            req_id, req = self.save_input_queue.get()
            for i in range(
                req.num_cpu_computed_tokens // self.block_size,
                min((req.num_computed_tokens + req.num_scheduled_tokens) // self.block_size, len(req.cpu_block_ids)),
            ):
                save_block_mapping.append((req.gpu_block_ids[i], req.cpu_block_ids[i]))
            with torch.npu.stream(self.save_stream):
                # MLA: kv_layer is tuple[tensor, tensor] means (rope, nope).
                # non-MLA: kv_layer is list[tensor], typically means [k, v].
                if self.use_mla:
                    start, step = self.tp_rank, self.tp_world_size
                else:
                    start, step = 0, 1
                for i in range(start, len(save_block_mapping), step):
                    gpu_block_id, cpu_block_id = save_block_mapping[i]
                    for cpu_kv_caches, gpu_kv_caches in zip(self.cpu_kv_caches, self.gpu_kv_caches.values()):
                        for cpu_layer_part, gpu_layer_part in zip(cpu_kv_caches, gpu_kv_caches):
                            cpu_layer_part[cpu_block_id].copy_(gpu_layer_part[gpu_block_id], non_blocking=True)
            self.save_stream.synchronize()
            self.save_output_queue.put(req_id)
            save_block_mapping.clear()


# copied and modified from vllm_ascend/worker/model_runner_v1.py
def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
    """
    Generates the KVCacheSpec by parsing the kv cache format from each
    Attention module in the static forward context.
    Returns:
        KVCacheSpec: A dictionary mapping layer names to their KV cache
        format. Layers that do not need KV cache are not included.
    """
    if has_ec_transfer() and get_ec_transfer().is_producer:
        return {}

    use_sparse = hasattr(vllm_config.model_config.hf_config, "index_topk")
    if vllm_config.cache_config.cache_dtype == "auto":
        kv_cache_dtype = vllm_config.model_config.dtype
    else:
        kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[vllm_config.cache_config.cache_dtype]

    kv_cache_spec: dict[str, KVCacheSpec] = {}
    attn_layers = get_layers_from_vllm_config(vllm_config, AttentionLayerBase)
    # NOTE: Must process Attention/MLAAttention before MambaBase to maintain
    # ordering expected by graph parameter update logic in attention backends.
    mamba_layers: dict[str, MambaBase] = {}
    for layer_name, attn_module in attn_layers.items():
        if isinstance(attn_module, Attention):
            if spec := attn_module.get_kv_cache_spec(vllm_config):
                kv_cache_spec[layer_name] = spec

        elif isinstance(attn_module, MLAAttention):
            if use_sparse:
                # TODO(cmq): This is a hack way to fix deepseek kvcache when
                # using DSA. Fix the spec in vLLM is the final way.
                block_size = vllm_config.cache_config.block_size
                kv_cache_spec[layer_name] = FullAttentionSpec(
                    block_size=block_size, num_kv_heads=1, head_size=attn_module.head_size, dtype=kv_cache_dtype
                )
            elif spec := attn_module.get_kv_cache_spec(vllm_config):
                kv_cache_spec[layer_name] = spec

        elif isinstance(attn_module, MambaBase):
            mamba_layers[layer_name] = attn_module

    if len(mamba_layers) > 0:
        if vllm_config.cache_config.enable_prefix_caching:
            raise NotImplementedError("Prefix caching is not supported for Mamba yet.")
        for layer_name, mamba_module in mamba_layers.items():
            if spec := mamba_module.get_kv_cache_spec(vllm_config):
                kv_cache_spec[layer_name] = spec

    return kv_cache_spec
