"""SnapshotLinkController — D→P RDMA snapshot pushes with dedicated GPU buffer.

Per `docs/SNAPSHOT_STORE_REFACTOR_ZH.md`, this controller now reserves a
dedicated GPU tensor (``snapshot_buf``) for receiving D→P snapshots, instead
of competing with the worker's ``token_to_kv_pool_allocator`` at
prepare_receive time. The kv_pool alloc is deferred to ``finalize_ingest``
when the bytes are already in hand — if that alloc fails we drop the
snapshot but RDMA reception itself succeeded.

Layout of the snapshot_buf for one session reception (chosen for
mooncake's batch_transfer_sync_write friendliness — every layer maps to
a single contiguous slab):

    [K_layer_0:  num_tokens × stride_k_bytes]
    [K_layer_1:  num_tokens × stride_k_bytes]
    ...
    [K_layer_L-1]
    [V_layer_0:  num_tokens × stride_v_bytes]
    ...
    [V_layer_L-1]

The buffer is split into multiple such slabs via ``SnapshotBufAllocator``.
"""

from __future__ import annotations

import logging
import os
import threading
import time
from dataclasses import dataclass, field
from typing import List, Optional, Tuple

logger = logging.getLogger(__name__)

# Env-var names (also exported from package __init__)
SNAPSHOT_LINK_ENABLE_ENV = "SGLANG_SNAPSHOT_LINK_ENABLE"
SNAPSHOT_LINK_HOST_ENV = "SGLANG_SNAPSHOT_LINK_HOST"
SNAPSHOT_LINK_PORT_ENV = "SGLANG_SNAPSHOT_LINK_PORT"
SNAPSHOT_LINK_IB_DEVICE_ENV = "SGLANG_SNAPSHOT_LINK_IB_DEVICE"

# Default snapshot_buf size: 8 GB. Enough for ~1.5 Qwen3-30B 50k-token sessions.
SNAPSHOT_BUF_BYTES_ENV = "SGLANG_SNAPSHOT_LINK_BUF_BYTES"
DEFAULT_SNAPSHOT_BUF_BYTES = 8 * 1024 * 1024 * 1024


@dataclass
class _LayerBufferDesc:
    """Per-layer KV buffer descriptor on this worker."""
    base_ptr: int             # data pointer of the layer's full buffer tensor
    bytes_per_token: int      # head_num * head_dim * dtype.itemsize
    capacity_bytes: int       # full buffer size in bytes
    is_k: bool                # True for K-buffer, False for V


@dataclass
class SnapshotIngestRecord:
    """P-side bookkeeping for one in-flight snapshot reception."""
    session_id: str
    slab_offset: int          # offset within snapshot_buf
    slab_size: int            # total bytes for this slab
    num_tokens: int
    k_layer_offsets: List[int]  # absolute byte offsets of K layers in snapshot_buf
    v_layer_offsets: List[int]
    per_token_k_bytes: int
    per_token_v_bytes: int
    created_at: float = field(default_factory=time.time)


class SnapshotBufAllocator:
    """First-fit free-list allocator over a single contiguous byte range.

    Tracks gaps in a sorted list. Merges adjacent free regions on free().
    """

    def __init__(self, capacity_bytes: int):
        self.capacity = capacity_bytes
        # Free regions sorted by offset: [(offset, size), ...]
        self._free: List[Tuple[int, int]] = [(0, capacity_bytes)]
        self._lock = threading.Lock()
        self._inflight: dict[int, int] = {}  # offset → size for sanity check

    def alloc(self, size: int) -> Optional[int]:
        """Return offset of allocated region, or None if no fit available."""
        if size <= 0:
            return None
        # Page-align allocations to 4 KB for RDMA-friendly alignment.
        size = (size + 4095) & ~4095
        with self._lock:
            for i, (off, sz) in enumerate(self._free):
                if sz >= size:
                    if sz == size:
                        self._free.pop(i)
                    else:
                        self._free[i] = (off + size, sz - size)
                    self._inflight[off] = size
                    return off
            return None

    def free(self, offset: int) -> bool:
        """Return True if the offset was successfully freed."""
        with self._lock:
            size = self._inflight.pop(offset, None)
            if size is None:
                return False
            # Insert sorted and merge adjacents
            self._free.append((offset, size))
            self._free.sort()
            merged: List[Tuple[int, int]] = []
            for off, sz in self._free:
                if merged and merged[-1][0] + merged[-1][1] == off:
                    merged[-1] = (merged[-1][0], merged[-1][1] + sz)
                else:
                    merged.append((off, sz))
            self._free = merged
            return True

    def available_bytes(self) -> int:
        with self._lock:
            return sum(sz for _, sz in self._free)

    def in_use_bytes(self) -> int:
        with self._lock:
            return sum(self._inflight.values())


def _import_transfer_engine():
    try:
        from mooncake.engine import TransferEngine
    except ImportError as e:
        raise ImportError(
            "mooncake.engine.TransferEngine is required for the snapshot "
            "link. Install mooncake-transfer-engine in the venv."
        ) from e
    return TransferEngine


class SnapshotLinkController:
    """Owns mooncake engine + kv_pool registrations + snapshot_buf + records.

    D-side use: push session KV via ``push_session_to_snapshot_buf``.
    P-side use: ``prepare_receive`` → caller pushes via RDMA →
                ``ingest_snapshot_into_kvpool`` (does GPU memcpy +
                radix insert) → ``finalize_record`` (frees the slab).
    """

    def __init__(
        self,
        host: str,
        port: int,
        ib_device: Optional[str],
        kv_pool_layer_buffers: List[Tuple[int, int, int, bool]],
        token_to_kv_pool_allocator,
        tree_cache=None,
        protocol: Optional[str] = None,
        snapshot_buf_bytes: Optional[int] = None,
    ):
        TransferEngine = _import_transfer_engine()
        self.host = host
        self.port = port
        self.ib_device = ib_device
        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
        self.tree_cache = tree_cache
        self.layer_buffers: List[_LayerBufferDesc] = [
            _LayerBufferDesc(
                base_ptr=base, bytes_per_token=btok,
                capacity_bytes=cap, is_k=is_k,
            )
            for (base, btok, cap, is_k) in kv_pool_layer_buffers
        ]

        self.engine = TransferEngine()
        proto = protocol or os.environ.get("MOONCAKE_PROTOCOL", "rdma")
        listen = f"{host}:{port}"
        ret = self.engine.initialize(listen, "P2PHANDSHAKE", proto, ib_device or "")
        if ret != 0:
            raise RuntimeError(
                f"SnapshotLinkController.initialize({listen}, {proto}, "
                f"ib={ib_device}) returned {ret}"
            )
        self._session_id = f"{host}:{self.engine.get_rpc_port()}"

        # Register existing kv_pool layer buffers (needed for D-side send and
        # for P-side ingest copy source = snapshot_buf, destination = kv_pool)
        ptrs = [d.base_ptr for d in self.layer_buffers]
        lens = [d.capacity_bytes for d in self.layer_buffers]
        try:
            reg_ret = self.engine.batch_register_memory(ptrs, lens)
        except Exception:
            reg_ret = 0
            for ptr, length in zip(ptrs, lens):
                r = self.engine.register_memory(ptr, length)
                if r != 0:
                    reg_ret = r
        if reg_ret != 0:
            logger.warning(
                "SnapshotLinkController kv_pool batch_register returned %d", reg_ret
            )

        # Allocate + register the dedicated snapshot reception buffer (P-side)
        # This decouples reception from kv_pool, avoiding the alloc-failed
        # death loop that killed E4-v4/v5.
        import torch

        if snapshot_buf_bytes is None:
            snapshot_buf_bytes = int(
                os.environ.get(SNAPSHOT_BUF_BYTES_ENV, DEFAULT_SNAPSHOT_BUF_BYTES)
            )
        device = self._allocator_device()
        try:
            self.snapshot_buf = torch.zeros(
                snapshot_buf_bytes, dtype=torch.uint8, device=device,
            )
        except RuntimeError as e:
            logger.warning(
                "Could not allocate snapshot_buf of %d bytes on %s: %s. "
                "Falling back to 1 GB.", snapshot_buf_bytes, device, e,
            )
            snapshot_buf_bytes = 1024 * 1024 * 1024
            self.snapshot_buf = torch.zeros(
                snapshot_buf_bytes, dtype=torch.uint8, device=device,
            )
        self._snapshot_buf_bytes = snapshot_buf_bytes
        self._snapshot_buf_ptr = self.snapshot_buf.data_ptr()
        ret = self.engine.register_memory(self._snapshot_buf_ptr, snapshot_buf_bytes)
        if ret != 0:
            logger.warning(
                "SnapshotLinkController snapshot_buf register_memory(%s, %d) ret=%d",
                hex(self._snapshot_buf_ptr), snapshot_buf_bytes, ret,
            )
        self.snapshot_buf_alloc = SnapshotBufAllocator(snapshot_buf_bytes)

        # Receive-side bookkeeping
        self._ingest_records: dict[str, SnapshotIngestRecord] = {}
        self._records_by_handle: dict[int, SnapshotIngestRecord] = {}
        self._next_handle = 1
        self._lock = threading.Lock()

        logger.info(
            "SnapshotLinkController up at %s (sid=%s, %d kv layer bufs, "
            "snapshot_buf=%.1f GB on %s)",
            listen, self._session_id, len(self.layer_buffers),
            snapshot_buf_bytes / 1e9, device,
        )

    # ----- accessors ----------------------------------------------------

    @property
    def snapshot_session_id(self) -> str:
        return self._session_id

    @property
    def snapshot_buf_ptr(self) -> int:
        return self._snapshot_buf_ptr

    @property
    def snapshot_buf_bytes(self) -> int:
        return self._snapshot_buf_bytes

    @property
    def layer_num(self) -> int:
        return len(self.layer_buffers) // 2

    def get_k_base_ptrs(self) -> List[int]:
        return [d.base_ptr for d in self.layer_buffers if d.is_k]

    def get_v_base_ptrs(self) -> List[int]:
        return [d.base_ptr for d in self.layer_buffers if not d.is_k]

    def get_stride_k_bytes(self) -> int:
        for d in self.layer_buffers:
            if d.is_k:
                return d.bytes_per_token
        return 0

    def get_stride_v_bytes(self) -> int:
        for d in self.layer_buffers:
            if not d.is_k:
                return d.bytes_per_token
        return 0

    def _allocator_device(self):
        # Best-effort: pull device from one of the buffer tensors via the allocator
        try:
            return self.token_to_kv_pool_allocator.device
        except AttributeError:
            return "cuda"

    # ----- P-side: prepare to receive ----------------------------------

    def prepare_receive(self, session_id: str, num_tokens: int) -> Optional[SnapshotIngestRecord]:
        """Carve a slab out of snapshot_buf large enough for num_tokens of K+V.

        Returns the record describing the slab layout, or None if snapshot_buf
        is full. This does NOT touch kv_pool — alloc happens at ingest time.
        """
        if num_tokens <= 0:
            return None
        stride_k = self.get_stride_k_bytes()
        stride_v = self.get_stride_v_bytes()
        L = self.layer_num
        slab_bytes = L * num_tokens * stride_k + L * num_tokens * stride_v
        offset = self.snapshot_buf_alloc.alloc(slab_bytes)
        if offset is None:
            logger.info(
                "prepare_receive: snapshot_buf full (sid=%s n=%d need=%d B available=%d B)",
                session_id, num_tokens, slab_bytes,
                self.snapshot_buf_alloc.available_bytes(),
            )
            return None
        # Layout: K0..KL-1, then V0..VL-1
        k_offs = [offset + i * num_tokens * stride_k for i in range(L)]
        v_offs = [offset + L * num_tokens * stride_k + i * num_tokens * stride_v
                  for i in range(L)]
        record = SnapshotIngestRecord(
            session_id=session_id,
            slab_offset=offset,
            slab_size=slab_bytes,
            num_tokens=num_tokens,
            k_layer_offsets=k_offs,
            v_layer_offsets=v_offs,
            per_token_k_bytes=stride_k,
            per_token_v_bytes=stride_v,
        )
        with self._lock:
            # Evict prior record for the same session (best-effort)
            old = self._ingest_records.pop(session_id, None)
            if old is not None:
                self.snapshot_buf_alloc.free(old.slab_offset)
                self._records_by_handle.pop(id(old), None)
            self._ingest_records[session_id] = record
            self._records_by_handle[id(record)] = record
        return record

    def lookup_by_handle(self, handle: int) -> Optional[SnapshotIngestRecord]:
        with self._lock:
            return self._records_by_handle.get(handle)

    def discard_record(self, session_id: str) -> None:
        with self._lock:
            rec = self._ingest_records.pop(session_id, None)
        if rec is not None:
            self.snapshot_buf_alloc.free(rec.slab_offset)
            with self._lock:
                self._records_by_handle.pop(id(rec), None)

    def total_pending_snapshot_bytes(self) -> int:
        with self._lock:
            return sum(rec.slab_size for rec in self._ingest_records.values())

    # ----- P-side: ingest snapshot into kv_pool + radix tree -----------

    def ingest_snapshot_into_kvpool(
        self,
        session_id: str,
        token_ids: List[int],
    ) -> Tuple[bool, str, int]:
        """Copy snapshot_buf bytes into kv_pool slots and insert into radix.

        Returns (ok, reason, inserted_prefix_len).
        """
        with self._lock:
            record = self._ingest_records.pop(session_id, None)
            if record is not None:
                self._records_by_handle.pop(id(record), None)
        if record is None:
            return False, "no-pending-ingest", 0

        try:
            n = min(len(token_ids), record.num_tokens)
            if n == 0:
                self.snapshot_buf_alloc.free(record.slab_offset)
                return False, "empty-token-ids", 0

            # Alloc kv_pool slots NOW that the snapshot bytes are in hand.
            try:
                indices_tensor = self.token_to_kv_pool_allocator.alloc(n)
            except Exception as exc:
                self.snapshot_buf_alloc.free(record.slab_offset)
                return False, f"kvpool-alloc-threw:{exc!r}", 0
            if indices_tensor is None:
                self.snapshot_buf_alloc.free(record.slab_offset)
                return False, "kvpool-alloc-failed-at-ingest", 0

            # GPU→GPU copy from snapshot_buf into kv_pool layer buffers
            try:
                self._copy_snapshot_to_kvpool(record, indices_tensor)
            except Exception as exc:
                logger.exception("snapshot→kvpool copy failed: %s", exc)
                # Free both allocations
                self._free_slot_indices(indices_tensor)
                self.snapshot_buf_alloc.free(record.slab_offset)
                return False, f"copy-failed:{exc!r}", 0

            # Insert into radix tree
            try:
                inserted_prefix_len = self._radix_insert(token_ids[:n], indices_tensor)
            except Exception as exc:
                logger.exception("radix insert failed: %s", exc)
                self._free_slot_indices(indices_tensor)
                self.snapshot_buf_alloc.free(record.slab_offset)
                return False, f"radix-insert-failed:{exc!r}", 0

            # Snapshot is now persisted into kv_pool + radix; the slab is no
            # longer needed.
            self.snapshot_buf_alloc.free(record.slab_offset)
            return True, "ok", int(inserted_prefix_len)
        except Exception as exc:
            # Belt-and-braces cleanup
            try:
                self.snapshot_buf_alloc.free(record.slab_offset)
            except Exception:
                pass
            return False, f"unexpected:{exc!r}", 0

    def _copy_snapshot_to_kvpool(
        self,
        record: SnapshotIngestRecord,
        slot_indices_tensor,
    ) -> None:
        """For each layer L: copy snapshot_buf[K_off[L]..] → k_buffer[L][slots]."""
        import torch

        n = record.num_tokens
        stride_k = record.per_token_k_bytes
        stride_v = record.per_token_v_bytes
        # View snapshot_buf as a 1-D byte tensor; slice by offsets.
        for L in range(self.layer_num):
            # K
            k_slab_start = record.k_layer_offsets[L] - record.slab_offset + record.slab_offset
            # NOTE: above is equivalent to record.k_layer_offsets[L] but kept for clarity
            k_slab_start = record.k_layer_offsets[L]
            k_layer_bytes = self.snapshot_buf[
                k_slab_start : k_slab_start + n * stride_k
            ].view(n, stride_k)
            # Compute destination tensor on kv_pool: dst[slot_indices] = src
            # We need access to the actual k_buffer[L] tensor. The controller
            # only has the raw ptr — so we materialize a view via from_blob-ish
            # trick. Easier: get the tensor from token_to_kv_pool_allocator's kvcache.
            kv_cache = self.token_to_kv_pool_allocator.get_kvcache()
            k_buf = kv_cache.k_buffer[L]  # (max_tokens, head, dim)
            # Flatten per-token to bytes
            flat = k_buf.view(k_buf.shape[0], -1)
            assert flat.shape[1] * flat.element_size() >= stride_k, (
                f"K layer {L} stride mismatch: pool {flat.shape[1] * flat.element_size()} vs snapshot {stride_k}"
            )
            # Copy: dst[slot_indices] ← src[:n]
            src_reshape = k_layer_bytes.view(n, flat.shape[1] * flat.element_size())
            # Byte-level view of destination rows
            dst_view = flat.view(torch.uint8)
            dst_view[slot_indices_tensor] = src_reshape

            # V
            v_slab_start = record.v_layer_offsets[L]
            v_layer_bytes = self.snapshot_buf[
                v_slab_start : v_slab_start + n * stride_v
            ]
            v_buf = kv_cache.v_buffer[L]
            v_flat = v_buf.view(v_buf.shape[0], -1)
            src_v = v_layer_bytes.view(n, v_flat.shape[1] * v_flat.element_size())
            v_dst_view = v_flat.view(torch.uint8)
            v_dst_view[slot_indices_tensor] = src_v

    def _radix_insert(self, token_ids: List[int], indices_tensor) -> int:
        """Insert (token_ids, kv_indices) into the underlying radix tree."""
        from sglang.srt.mem_cache.base_prefix_cache import InsertParams
        from sglang.srt.mem_cache.radix_cache import RadixKey
        from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache

        inner = self.tree_cache
        if isinstance(inner, SessionAwareCache):
            inner = inner.inner
        if inner is None:
            raise RuntimeError("tree_cache not provided to SnapshotLinkController")
        radix_key = RadixKey(token_ids, None)
        result = inner.insert(InsertParams(key=radix_key, value=indices_tensor))
        return int(getattr(result, "prefix_len", 0))

    def _free_slot_indices(self, indices_tensor) -> None:
        try:
            self.token_to_kv_pool_allocator.free(indices_tensor)
        except Exception as e:
            logger.warning("_free_slot_indices failed: %s", e)

    # ----- D-side: push session KV to a peer's snapshot_buf ------------

    def push_session_to_snapshot_buf(
        self,
        *,
        target_snapshot_session_id: str,
        src_slot_indices: List[int],
        target_snapshot_buf_base: int,
        target_k_layer_offsets: List[int],
        target_v_layer_offsets: List[int],
        target_per_token_k_bytes: int,
        target_per_token_v_bytes: int,
    ) -> Tuple[int, int]:
        """Push session KV from local kv_pool into a peer's snapshot_buf slab.

        For each layer: gather src ranges (possibly scattered slot indices)
        and write to a contiguous range in the peer's snapshot_buf.
        Returns (mooncake_return_code, bytes_pushed).
        """
        if not src_slot_indices:
            return 0, 0
        layer_num = self.layer_num
        k_src_bases = self.get_k_base_ptrs()
        v_src_bases = self.get_v_base_ptrs()
        stride_k = self.get_stride_k_bytes()
        stride_v = self.get_stride_v_bytes()
        if (len(target_k_layer_offsets) != layer_num
                or len(target_v_layer_offsets) != layer_num):
            raise ValueError(
                f"target K/V layer offset count {len(target_k_layer_offsets)}/"
                f"{len(target_v_layer_offsets)} != local layer_num {layer_num}"
            )
        if (stride_k != target_per_token_k_bytes
                or stride_v != target_per_token_v_bytes):
            raise ValueError(
                f"stride mismatch: local k={stride_k}/v={stride_v}, "
                f"target k={target_per_token_k_bytes}/v={target_per_token_v_bytes}"
            )
        n = len(src_slot_indices)

        local_addrs: List[int] = []
        remote_addrs: List[int] = []
        lengths: List[int] = []

        # Coalesce contiguous src runs.
        # Inner-loop helper to walk indices and emit run boundaries.
        def _emit_runs(src_base: int, tgt_base: int, stride: int) -> None:
            run_src_start = run_tgt_start = run_len = None
            for tgt_idx, src in enumerate(src_slot_indices):
                if run_src_start is None:
                    run_src_start, run_tgt_start, run_len = src, tgt_idx, 1
                elif src == run_src_start + run_len:
                    run_len += 1
                else:
                    local_addrs.append(src_base + run_src_start * stride)
                    remote_addrs.append(tgt_base + run_tgt_start * stride)
                    lengths.append(run_len * stride)
                    run_src_start, run_tgt_start, run_len = src, tgt_idx, 1
            if run_src_start is not None:
                local_addrs.append(src_base + run_src_start * stride)
                remote_addrs.append(tgt_base + run_tgt_start * stride)
                lengths.append(run_len * stride)

        for L in range(layer_num):
            _emit_runs(
                k_src_bases[L],
                target_snapshot_buf_base + target_k_layer_offsets[L],
                stride_k,
            )
            _emit_runs(
                v_src_bases[L],
                target_snapshot_buf_base + target_v_layer_offsets[L],
                stride_v,
            )

        t0 = time.perf_counter()
        try:
            ret = self.engine.batch_transfer_sync_write(
                target_snapshot_session_id, local_addrs, remote_addrs, lengths,
            )
        except Exception as e:
            logger.exception(
                "SnapshotLinkController.push_session_to_snapshot_buf threw: %s", e
            )
            return -1, 0
        t1 = time.perf_counter()
        bytes_pushed = sum(lengths)
        logger.info(
            "push_session_to_snapshot_buf → %s: %d ops, %d B, ret=%d, %.2f ms",
            target_snapshot_session_id, len(lengths), bytes_pushed, ret,
            (t1 - t0) * 1000.0,
        )
        return ret, bytes_pushed
