# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import threading
import time
from dataclasses import dataclass

import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

from vllm.config import VllmConfig
from vllm.distributed.kv_transfer.kv_connector.utils import EngineId
from vllm.logger import init_logger

WorkerAddr = str

logger = init_logger(__name__)


class RegisterWorkerPayload(BaseModel):
    engine_id: EngineId
    dp_rank: int
    tp_rank: int
    pp_rank: int
    addr: WorkerAddr


@dataclass
class EngineEntry:
    engine_id: EngineId
    # {tp_rank: {pp_rank: worker_addr}}
    worker_addr: dict[int, dict[int, WorkerAddr]]


class QueryBlocksRequest(BaseModel):
    block_hashes: list[str] | None = None  # hex-encoded BlockHash values (legacy)
    token_ids: list[int] | None = None     # raw token IDs for C-side hash lookup
    num_tokens: int | None = None          # number of tokens to match prefix for
    pin_token: str


class QueryBlocksResponse(BaseModel):
    block_ids: list[int | None]  # None = cache miss (prefix match stops)
    kv_caches_base_addr: list[int]
    block_len: int
    hostname: str
    rpc_port: int


class PushBlocksRequest(BaseModel):
    token_ids: list[int]
    num_tokens: int
    dst_block_ids: list[int]  # D's allocated block IDs for receiving
    dst_base_addrs: list[int]  # D's kv_caches_base_addr
    dst_block_len: int  # D's block_len
    dst_session: str  # D's "hostname:rpc_port" for RDMA write


class PushBlocksResponse(BaseModel):
    matched: int
    pushed: bool


class EstimateHitRequest(BaseModel):
    token_ids: list[int]
    block_size: int = 512


class EstimateHitResponse(BaseModel):
    hit_tokens: int


class UnpinBlocksRequest(BaseModel):
    pin_token: str


class MooncakeBootstrapServer:
    """
    A centralized server running on the global rank 0 prefiller worker.
    Prefiller workers register their connection info (IP, port, ranks) here.
    Also serves block mapping queries for direct RDMA read.
    """

    def __init__(self, vllm_config: VllmConfig, host: str, port: int):
        self.workers: dict[int, EngineEntry] = {}

        self.host = host
        self.port = port
        self.app = FastAPI()
        self._register_routes()
        self.server_thread: threading.Thread | None = None
        self.server: uvicorn.Server | None = None

        # Direct RDMA read support
        self._hash_table: dict[str, int] = {}  # hex BlockHash → block_id
        self._token_hash_table: dict[str, int] = {}  # str(hash(token_tuple)) → block_id
        self._kv_info: dict | None = None
        self._pinned: dict[str, list[int]] = {}
        self._block_pool = None

    def __del__(self):
        self.shutdown()

    def _register_routes(self):
        self.app.post("/register")(self.register_worker)
        self.app.get("/query", response_model=dict[int, EngineEntry])(self.query)
        self.app.post("/query_blocks")(self.query_blocks)
        self.app.post("/unpin_blocks")(self.unpin_blocks)
        self.app.post("/push_blocks")(self.push_blocks)
        self.app.post("/estimate_hit")(self.estimate_hit)

    def start(self):
        if self.server_thread:
            return

        config = uvicorn.Config(app=self.app, host=self.host, port=self.port)
        self.server = uvicorn.Server(config=config)
        self.server_thread = threading.Thread(
            target=self.server.run, name="mooncake_bootstrap_server", daemon=True
        )
        self.server_thread.start()
        while not self.server.started:
            time.sleep(0.1)  # Wait for the server to start
        logger.info("Mooncake Bootstrap Server started at %s:%d", self.host, self.port)

    def shutdown(self):
        if self.server_thread is None or self.server is None or not self.server.started:
            return

        self.server.should_exit = True
        self.server_thread.join()
        logger.info("Mooncake Bootstrap Server stopped.")

    async def register_worker(self, payload: RegisterWorkerPayload):
        """Handles registration of a prefiller worker."""
        if payload.dp_rank not in self.workers:
            self.workers[payload.dp_rank] = EngineEntry(
                engine_id=payload.engine_id,
                worker_addr={},
            )

        dp_entry = self.workers[payload.dp_rank]
        if dp_entry.engine_id != payload.engine_id:
            raise HTTPException(
                status_code=400,
                detail=(
                    f"Engine ID mismatch for dp_rank={payload.dp_rank}: "
                    f"expected {dp_entry.engine_id}, got {payload.engine_id}"
                ),
            )
        if payload.tp_rank not in dp_entry.worker_addr:
            dp_entry.worker_addr[payload.tp_rank] = {}

        tp_entry = dp_entry.worker_addr[payload.tp_rank]
        if payload.pp_rank in tp_entry:
            raise HTTPException(
                status_code=400,
                detail=(
                    f"Worker with dp_rank={payload.dp_rank}, "
                    f"tp_rank={payload.tp_rank}, pp_rank={payload.pp_rank} "
                    f"is already registered at "
                    f"{tp_entry[payload.pp_rank]}, "
                    f"but still want to register at {payload.addr}"
                ),
            )

        tp_entry[payload.pp_rank] = payload.addr
        logger.debug(
            "Registered worker: engine_id=%s, dp_rank=%d, tp_rank=%d, pp_rank=%d at %s",
            payload.engine_id,
            payload.dp_rank,
            payload.tp_rank,
            payload.pp_rank,
            payload.addr,
        )

        return {"status": "ok"}

    async def query(self) -> dict[int, EngineEntry]:
        return self.workers

    def set_worker_kv_info(
        self,
        kv_caches_base_addr: list[int],
        block_len: int,
        block_size: int,
        hostname: str,
        rpc_port: int,
        transfer_engine=None,
    ):
        self._kv_info = {
            "kv_caches_base_addr": kv_caches_base_addr,
            "block_len": block_len,
            "block_size": block_size,
            "hostname": hostname,
            "rpc_port": rpc_port,
        }
        self._transfer_engine = transfer_engine

    def update_hash_table(
        self,
        updates: dict[str, int],
        removals: set[str],
    ):
        for k in removals:
            self._hash_table.pop(k, None)
        self._hash_table.update(updates)

    def set_block_pool(self, block_pool):
        """Store reference to scheduler's block pool for token-based lookup."""
        self._block_pool = block_pool

    async def query_blocks(self, req: QueryBlocksRequest):
        if self._kv_info is None:
            raise HTTPException(503, "Worker KV info not registered yet")

        block_ids: list[int | None] = []
        pinned_ids: list[int] = []

        if req.token_ids is not None and self._hash_table:
            # Token-based lookup: compute hashes from tokens, match against synced hash table
            block_ids, pinned_ids = self._lookup_by_tokens(
                req.token_ids, req.num_tokens)
        elif req.block_hashes is not None:
            # Hash-based lookup (legacy)
            for h in req.block_hashes:
                bid = self._hash_table.get(h)
                if bid is not None:
                    block_ids.append(bid)
                    pinned_ids.append(bid)
                else:
                    block_ids.append(None)
                    break

        logger.info(
            "query_blocks: %d/%d matched (token_mode=%s, hash_table=%d)",
            len(pinned_ids),
            req.num_tokens // self._kv_info.get("block_size", 512)
            if req.num_tokens else len(req.block_hashes or []),
            req.token_ids is not None,
            len(self._hash_table),
        )

        self._pinned[req.pin_token] = pinned_ids
        return QueryBlocksResponse(
            block_ids=block_ids,
            kv_caches_base_addr=self._kv_info["kv_caches_base_addr"],
            block_len=self._kv_info["block_len"],
            hostname=self._kv_info["hostname"],
            rpc_port=self._kv_info["rpc_port"],
        )

    def _lookup_by_tokens(
        self, token_ids: list[int], num_tokens: int | None
    ) -> tuple[list[int | None], list[int]]:
        """Look up cached blocks by computing hashes using the synced hash table.

        Uses module-level NONE_HASH (accessed via module ref to get latest value
        after init_none_hash is called).
        """
        import vllm.v1.core.kv_cache_utils as kv_utils
        from vllm.utils.hashing import sha256

        block_size = self._kv_info.get("block_size", 512) if self._kv_info else 512
        n = num_tokens or len(token_ids)
        n = min(n, len(token_ids))
        num_blocks = n // block_size

        block_ids: list[int | None] = []
        pinned_ids: list[int] = []
        prev_hash = kv_utils.NONE_HASH  # module-level ref, always current

        for i in range(num_blocks):
            block_tokens = tuple(token_ids[i * block_size:(i + 1) * block_size])
            block_hash = kv_utils.hash_block_tokens(
                sha256, prev_hash, block_tokens, None)
            prev_hash = block_hash

            bid = self._hash_table.get(block_hash.hex())
            if i == 0:
                table_sample = next(iter(self._hash_table)) if self._hash_table else "empty"
                logger.info(
                    "_lookup: hash=%s NONE=%s tbl=%s",
                    block_hash.hex()[:12], kv_utils.NONE_HASH.hex()[:12], table_sample[:12])
            if bid is not None:
                block_ids.append(bid)
                pinned_ids.append(bid)
            else:
                if i == 0:
                    block_ids.append(None)
                break

        return block_ids, pinned_ids

    async def unpin_blocks(self, req: UnpinBlocksRequest):
        self._pinned.pop(req.pin_token, None)
        return {"status": "ok"}

    async def estimate_hit(self, req: EstimateHitRequest):
        """Read-only probe: how many prefix-contiguous tokens are cached?

        Reuses _lookup_by_tokens (proven to work with push_blocks) instead
        of reimplementing hash computation.
        """
        if self._kv_info is None:
            raise HTTPException(503, "Worker KV info not registered yet")

        if not self._hash_table:
            return EstimateHitResponse(hit_tokens=0)

        block_ids, _ = self._lookup_by_tokens(req.token_ids, None)
        hit_blocks = sum(1 for b in block_ids if b is not None)
        block_size = self._kv_info.get("block_size", 512)
        hit_tokens = hit_blocks * block_size

        logger.info("estimate_hit: %d/%d blocks hit (%d tokens, tbl=%d)",
                    hit_blocks, len(req.token_ids) // block_size,
                    hit_tokens, len(self._hash_table))
        return EstimateHitResponse(hit_tokens=hit_tokens)

    async def push_blocks(self, req: PushBlocksRequest):
        """Query matching blocks by token_ids, then PUSH them to D via RDMA write."""
        if self._kv_info is None or self._transfer_engine is None:
            raise HTTPException(503, "Worker not ready")

        block_ids, _ = self._lookup_by_tokens(req.token_ids, req.num_tokens)
        matched_src = [b for b in block_ids if b is not None]
        num_matched = len(matched_src)

        if num_matched == 0:
            logger.info("push_blocks: 0 matched")
            return PushBlocksResponse(matched=0, pushed=False)

        matched_dst = req.dst_block_ids[:num_matched]
        src_base = self._kv_info["kv_caches_base_addr"]
        src_block_len = self._kv_info["block_len"]

        src_ptrs: list[int] = []
        dst_ptrs: list[int] = []
        lengths: list[int] = []

        for src_layer, dst_layer in zip(src_base, req.dst_base_addrs):
            for s_bid, d_bid in zip(matched_src, matched_dst):
                src_ptrs.append(src_layer + s_bid * src_block_len)
                dst_ptrs.append(dst_layer + d_bid * req.dst_block_len)
                lengths.append(src_block_len)

        import asyncio
        loop = asyncio.get_event_loop()
        ret = await loop.run_in_executor(
            None,
            self._transfer_engine.batch_transfer_sync_write,
            req.dst_session, src_ptrs, dst_ptrs, lengths,
        )

        logger.info("push_blocks: %d matched, push ret=%d", num_matched, ret)
        return PushBlocksResponse(matched=num_matched, pushed=(ret == 0))
