# Copyright 2026 The HuggingFace Inc. team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Centralized offloading logic for continuous batching.

Handles two offloading strategies when the GPU KV cache is full:
  1. CPU offloading: copy the KV cache to a pre-allocated pinned CPU buffer, preserving exact request state.
  2. Soft reset: discard the KV cache and re-prefill from scratch when the request is re-scheduled. This incurs no data
    transfer overhead, but we need to re-run prefill over all initial + generated tokens (so more compute overhead).

The CPU swap pool is a static set of pinned tensors allocated once at init (like vLLM/SGLang). Blocks are tracked
with a simple free set — no dynamic allocation or deallocation of tensors ever happens at runtime.
"""

import logging
from contextlib import nullcontext
from itertools import chain

import torch

from ...utils import is_psutil_available
from .cache import PagedAttentionCache
from .distributed import DistributedHelper
from .requests import FutureRequestState, RequestState, RequestStatus, logger
from .scheduler import Scheduler


def contiguous_runs(indices: list[int]) -> list[tuple[int, int, int]]:
    """Groups an index list into (start_index, offset, length) runs of consecutive values, so scattered block copies
    can be performed as a few slice copies."""
    runs = []
    if not indices:
        return runs
    start = prev = indices[0]
    offset = 0
    for i in range(1, len(indices)):
        if indices[i] != prev + 1:
            runs.append((start, offset, i - offset))
            start, offset = indices[i], i
        prev = indices[i]
    runs.append((start, offset, len(indices) - offset))
    return runs


class OffloadingManager:
    """Manages request offloading and restoration for continuous batching.

    Owns a static CPU swap pool (pre-allocated pinned tensors mirroring the GPU cache layout), performs GPU↔CPU block
    copies, decides between CPU offloading and soft reset, and ensures cleanup on cancellation/failure/reset.
    """

    def __init__(
        self,
        cache: PagedAttentionCache,
        scheduler: Scheduler,
        cpu_offload_space_gib: float | None,
        safety_threshold: float,
        compute_stream: torch.cuda.Stream | None,
        distributed_helper: DistributedHelper,
    ) -> None:
        self.cache = cache
        self.scheduler = scheduler
        # All offloading transfers run on the compute stream (stream-ordered, like the fork copy path)
        self._compute_stream = compute_stream

        # Bookkeeping defaults, valid whether or not the pool is allocated
        self._cpu_key_cache: list[torch.Tensor] = []
        self._cpu_value_cache: list[torch.Tensor] = []
        self._gpu_key_views: list[torch.Tensor] = []
        self._gpu_value_views: list[torch.Tensor] = []
        self._free_cpu_blocks: list[int] = []
        self._request_id_to_cpu_blocks: dict[str, list[int]] = {}
        self._request_id_to_group_block_counts: dict[str, list[int]] = {}

        # Compute the size of the CPU swap pool in blocks
        num_cpu_blocks = self._compute_num_cpu_blocks(cpu_offload_space_gib, safety_threshold)
        num_cpu_blocks = torch.tensor(num_cpu_blocks, dtype=torch.int32, device="cpu")
        self._num_cpu_blocks = int(distributed_helper.tp_all_reduce_min(num_cpu_blocks, on_cpu=True).item())

        offloading_enabled = cpu_offload_space_gib is not None and cpu_offload_space_gib > 0
        if self._num_cpu_blocks == 0:
            if offloading_enabled:
                logger.warning(
                    f"cpu_offload_space={cpu_offload_space_gib:.1f} GiB is too small for even one block. "
                    "No CPU offloading."
                )
            return None

        # Allocate the CPU swap pool
        cpu_cache_shape = (self._num_cpu_blocks, cache.block_size, cache.num_key_value_heads, cache.head_dim)
        for _ in cache.key_cache:
            self._cpu_key_cache.append(torch.empty(cpu_cache_shape, dtype=cache.dtype, pin_memory=True))
            self._cpu_value_cache.append(torch.empty(cpu_cache_shape, dtype=cache.dtype, pin_memory=True))

        # Pre-view the GPU cache tensors as block-shaped so the hot copy paths avoid per-op .view() calls
        block_shape = (-1, cache.block_size, cache.num_key_value_heads, cache.head_dim)
        for k_cache, v_cache in zip(cache.key_cache, cache.value_cache):
            self._gpu_key_views.append(k_cache.view(*block_shape))
            self._gpu_value_views.append(v_cache.view(*block_shape))

        # The free list is kept sorted so bulk offloads land in (mostly) contiguous pool slices
        self._free_cpu_blocks = list(range(self._num_cpu_blocks))

        # Log the size of the CPU swap pool
        cache_tensor = self._cpu_key_cache[0]
        size_in_bytes = 2 * cache_tensor.numel() * cache_tensor.element_size() * len(cache.key_cache)
        logger.info(
            f"CPU swap pool initialized: {self._num_cpu_blocks} blocks ({size_in_bytes / (1024**3):.2f} GiB pinned)"
        )

    def _compute_num_cpu_blocks(self, cpu_offload_space_gib: float | None, safety_threshold: float) -> int:
        """Returns the number of blocks that can fit in the CPU swap pool."""
        # Compute the CPU pool size in bytes
        offload_bytes = int(cpu_offload_space_gib * (1024**3)) if cpu_offload_space_gib is not None else None

        # Determine the maximum number of bytes that can be offloaded based on the safety threshold
        if is_psutil_available():
            import psutil

            total_ram = psutil.virtual_memory().available
            max_bytes = int(total_ram * safety_threshold)
        else:
            max_bytes = None

        # If both the request number of bytes and its limit are not None, we just clamp one to the other
        if offload_bytes is not None and max_bytes is not None:
            if offload_bytes > max_bytes:
                clamped_gib = max_bytes / (1024**3)
                logger.warning(
                    f"cpu_offload_space={cpu_offload_space_gib:.1f} GiB exceeds {safety_threshold:.0%} of total RAM "
                    f"({total_ram / (1024**3):.1f} GiB). Clamping to {clamped_gib:.1f} GiB."
                )
                offload_bytes = max_bytes
        # Else if the max is None, throw a warning and accept the requested number of bytes as is
        elif offload_bytes is not None:
            logger.warning(
                "psutil is not available — cpu_offload_space_safety_threshold cannot be enforced. "
                "Install psutil to enable the safety cap."
            )
        # Else if the requested number of bytes is None, we use the max number of bytes as the requested number of bytes
        elif max_bytes is not None:
            offload_bytes = max_bytes
            logger.warning(f"Auto-sizing CPU swap pool from safety threshold: {max_bytes / (1024**3):.2f} GiB.")
        # Otherwise, it means the pool was supposed to be sized using psutil but it is not available
        else:
            raise ImportError(
                "cpu_offload_space=None requires psutil to auto-size the CPU swap pool. Install psutil or pass an "
                "explicit GiB value."
            )

        # Compute how many blocks fit in CPU pool
        bytes_per_block = (
            2                                 # one for key, one for value
            * len(self.cache.key_cache)       # number of layers in a layer group
            * self.cache.block_size           # block size
            * self.cache.num_key_value_heads  # number of key value heads
            * self.cache.head_dim             # head dimension
            * self.cache.dtype.itemsize       # data type size in bytes
        )  # fmt: skip
        if bytes_per_block == 0:
            raise ValueError("The number of bytes per block is 0. This is not possible.")
        return offload_bytes // bytes_per_block

    def _stream_ctx(self):
        """Returns a context manager that runs enclosed ops on the compute stream, or a no-op when none is set."""
        return torch.cuda.stream(self._compute_stream) if self._compute_stream is not None else nullcontext()

    def offload_requests(self) -> int:
        """Evict enough active requests that, at the next batch, every remaining starved request can allocate the
        blocks it needs. Victims are taken from the starved requests reported by the scheduler, newest first, so the
        batch that was just scheduled is never touched and the demand directly bounds the amount of offloading.
        Tries CPU offloading first; victims that do not fit in the pool are soft reset. Returns the number of evicted
        requests."""
        scheduler = self.scheduler
        starved = scheduler.starved_requests
        if not starved:
            return 0

        # Pick victims until the demand of the remaining starved requests fits in the free blocks. Evicting a victim
        # both removes its own demand and frees its blocks. We never evict the last active request.
        free_blocks = self.cache.get_num_free_blocks()
        demand = sum(blocks_needed for _, blocks_needed in starved)
        num_active = len(scheduler.active_requests)
        victims: list[RequestState] = []
        while demand > free_blocks and starved and num_active - len(victims) > 1:
            state, blocks_needed = starved.pop()
            victims.append(state)
            demand -= blocks_needed
            free_blocks += self.cache.blocks_in_use(state.request_id)  # approximation because of prefix sharing
        if not victims:
            return 0
        victims.reverse()  # ensures the oldest (ie. largest) requests are offloaded first

        # Copy as many victims as fit in the CPU pool, in one batched copy. Must happen before the blocks are freed.
        cpu_offloaded = self._offload_to_cpu(victims)
        # Requeue victims oldest-first so they will become active again in (roughly) their original order
        for state in victims:
            request_id = state.request_id
            if request_id in cpu_offloaded:
                # We set the allocated blocks to 0 so the scheduler re-allocates all blocks using position_offset.
                state.allocated_blocks = 0
                if state._status == RequestStatus.DECODING:
                    # In async mode, a request can be offloaded for preparation of batch N+1 while still in flight in
                    # batch N. Since the token generated by batch N will be discarded at the update, we roll back one
                    # token to avoid restoring with fake information (placeholder token or partial KV).
                    if state.position_offset == len(state.initial_tokens) + len(state.generated_tokens):
                        state.position_offset -= 1
                        last_true_token = (state.generated_tokens or state.initial_tokens)[-1]
                        state.remaining_prefill_tokens = [last_true_token]
                    # Otherwise the next token is known: re-processing it on restore continues the request exactly.
                    else:
                        state.remaining_prefill_tokens = state.tokens_to_process[:]
                # The new state is the same as the old one, but with the status set to PENDING. We bypass the setter
                # to avoid the lifespan bookkeeping and the associated warning
                state._status = RequestStatus.PENDING
                new_state = state
            else:
                new_state = state.create_equivalent_initial_request()
                state._status = RequestStatus.FINISHED
            scheduler.finish_request(request_id)
            scheduler.add_waiting_request(new_state)

        scheduler.block_new_requests = True
        if logger.isEnabledFor(logging.INFO):
            free_blocks = self.cache.get_num_free_blocks()
            logger.info(
                f"Offloaded {len(victims)} requests ({len(cpu_offloaded)} to CPU, {len(victims) - len(cpu_offloaded)} "
                f"soft reset): {len(starved)} starved requests remain for {free_blocks} free blocks."
            )
        return len(victims)

    def restore_scheduled_requests(self, requests_in_batch: list[FutureRequestState]) -> None:
        """Restore KV caches from CPU for any CPU-offloaded requests in the scheduled batch. Indices are accumulated
        per group across all requests, then copied in one batched operation per layer."""
        cache = self.cache
        all_cpu_indices: list[int] = []
        all_gpu_indices: list[int] = []

        for future_state in requests_in_batch:
            # Skip state that are not CPU-offloaded
            state = future_state.state
            if not state.is_cpu_offloaded:
                continue
            # TODO: if the H2D copy below raises, already-popped entries leak (never returned to _free_cpu_blocks)
            # Accumulate CPU indices for this request
            cpu_indices = self._request_id_to_cpu_blocks.pop(state.request_id)
            group_counts = self._request_id_to_group_block_counts.pop(state.request_id)
            all_cpu_indices.extend(cpu_indices)
            # Accumulate GPU indices for this request, but since there may be extra block due to re-allocation, slice to
            # match the number of blocks offloaded.
            max_allocated_blocks = 0
            for group_idx, n in enumerate(group_counts):
                gpu_blocks = cache.group_cache_managers[group_idx].block_table.get(state.request_id, [])
                all_gpu_indices.extend(gpu_blocks[:n])
                # The allocated count must reflect the blocks held NOW (after re-allocation), not the offloaded count:
                # undercounting creates phantom block demand that can starve the request forever.
                max_allocated_blocks = max(max_allocated_blocks, len(gpu_blocks))
            # Restore the state to non-offloaded state
            state.is_cpu_offloaded = False
            state.allocated_blocks = max_allocated_blocks
            # Prefix sharing: restored blocks will be re-hashed during the next update
            if cache.allow_block_sharing:
                future_state.complete_blocks += state.position_offset // cache.block_size
            logger.debug(
                f"Restored CPU-offloaded request {state.request_id} with {len(state.initial_tokens)} prefill tokens "
                f"and {len(state.generated_tokens)} generated tokens."
            )

        # Early return if there are no copy to perform
        if not all_cpu_indices:
            return None

        # Single batched copy for all requests: a few non-blocking slice copies into a staging tensor per layer, then
        # one scatter into the cache. All stream-ordered, so the host never waits and the next forward sees the data.
        n = len(all_cpu_indices)
        runs = contiguous_runs(all_cpu_indices)
        cache = self.cache
        with self._stream_ctx():
            gpu_ids = torch.as_tensor(all_gpu_indices, dtype=torch.long).to(cache.device, non_blocking=True)
            staging_shape = (n, cache.block_size, cache.num_key_value_heads, cache.head_dim)
            staging = torch.empty(staging_shape, dtype=cache.dtype, device=cache.device)
            cpu_caches = chain(self._cpu_key_cache, self._cpu_value_cache)
            gpu_views = chain(self._gpu_key_views, self._gpu_value_views)
            for cpu_cache, gpu_view in zip(cpu_caches, gpu_views):
                for start, offset, length in runs:
                    staging[offset : offset + length].copy_(cpu_cache[start : start + length], non_blocking=True)
                gpu_view.index_copy_(0, gpu_ids, staging)
        self._free_cpu_blocks = sorted(self._free_cpu_blocks + all_cpu_indices)

    def free_request_cpu_cache(self, state: RequestState, keep_unsorted: bool = False) -> None:
        """Free CPU blocks for a single request (e.g., on cancellation)."""
        if state.is_cpu_offloaded:
            self._return_cpu_blocks(state.request_id)
            state.is_cpu_offloaded = False
            if not keep_unsorted:
                self._free_cpu_blocks.sort()

    def free_all_waiting_cpu_caches(self) -> None:
        """Free all CPU-offloaded caches in the waiting queue (e.g., on fail_all or reset)."""
        for state in self.scheduler.waiting_requests.values():
            self.free_request_cpu_cache(state, keep_unsorted=True)
        self._free_cpu_blocks.sort()

    def reset(self) -> None:
        """Reset CPU offloading state for a new generation session."""
        self.free_all_waiting_cpu_caches()
        self._request_id_to_cpu_blocks.clear()
        self._request_id_to_group_block_counts.clear()
        self._free_cpu_blocks = list(range(self._num_cpu_blocks))

    def _offload_to_cpu(self, victims: list[RequestState]) -> set[str]:
        """Copy the KV cache blocks of as many victims as fit in the CPU swap pool from GPU to the pool, in one
        batched, non-blocking copy per layer. Returns the request ids that were offloaded.

        All transfers are enqueued on the compute stream with pinned destinations, so the host never waits on them:
        correctness is guaranteed by stream ordering, since restores and cache writes go through the same stream.
        """
        # Select the victims that fit in the pool and gather their GPU block indices
        offloaded: list[tuple[RequestState, list[int], list[int]]] = []
        all_gpu_indices: list[int] = []
        free_pool_blocks = len(self._free_cpu_blocks)
        for state in victims:
            gpu_indices = []
            group_block_counts = []
            for cm in self.cache.group_cache_managers:
                blocks = cm.block_table.get(state.request_id, [])
                gpu_indices.extend(blocks)
                group_block_counts.append(len(blocks))
            # If there is enough free CPU blocks, offload the request to CPU
            if gpu_indices and len(all_gpu_indices) + len(gpu_indices) <= free_pool_blocks:
                offloaded.append((state, gpu_indices, group_block_counts))
                all_gpu_indices.extend(gpu_indices)

        # If no requests were offloaded to CPU, we can stop here
        if not offloaded:
            return set()

        # Reserve the smallest free CPU blocks: the free list is sorted, so destinations form few contiguous runs
        n = len(all_gpu_indices)
        all_cpu_indices = self._free_cpu_blocks[:n]
        self._free_cpu_blocks = self._free_cpu_blocks[n:]
        runs = contiguous_runs(all_cpu_indices)

        # One gather and a few slice copies per layer, all stream-ordered and non-blocking for the host
        with self._stream_ctx():
            gpu_ids = torch.as_tensor(all_gpu_indices, dtype=torch.long).to(self.cache.device, non_blocking=True)
            gpu_views = chain(self._gpu_key_views, self._gpu_value_views)
            cpu_caches = chain(self._cpu_key_cache, self._cpu_value_cache)
            for gpu_view, cpu_cache in zip(gpu_views, cpu_caches):
                gathered_blocks = gpu_view.index_select(0, gpu_ids)
                for start, offset, length in runs:
                    cpu_cache[start : start + length].copy_(
                        gathered_blocks[offset : offset + length], non_blocking=True
                    )

        # No explicit sync needed: finish_request is logical, and the next forward pass serializes on the same stream
        offset = 0
        for state, gpu_indices, group_block_counts in offloaded:
            self._request_id_to_cpu_blocks[state.request_id] = all_cpu_indices[offset : offset + len(gpu_indices)]
            self._request_id_to_group_block_counts[state.request_id] = group_block_counts
            state.is_cpu_offloaded = True
            offset += len(gpu_indices)
        return {state.request_id for state, _, _ in offloaded}

    def _return_cpu_blocks(self, request_id: str) -> tuple[list[int], list[int]]:
        """Return CPU blocks to the free pool without copying anything."""
        cpu_ids = self._request_id_to_cpu_blocks.pop(request_id)
        group_counts = self._request_id_to_group_block_counts.pop(request_id)
        self._free_cpu_blocks.extend(cpu_ids)
        return cpu_ids, group_counts
