# SPDX-License-Identifier: Apache-2.0
# Standard
from typing import TYPE_CHECKING, Optional
import asyncio
import threading

# Third Party
import msgspec
import zmq
import zmq.asyncio

# First Party
from lmcache.logging import init_logger
from lmcache.v1.cache_controller.full_sync_sender import FullSyncSender
from lmcache.v1.cache_controller.message import (
    BatchedP2PLookupMsg,
    BatchedP2PLookupRetMsg,
    ClearWorkerMsg,
    ClearWorkerRetMsg,
    CompressWorkerMsg,
    CompressWorkerRetMsg,
    DecompressWorkerMsg,
    DecompressWorkerRetMsg,
    DeRegisterMsg,
    ErrorMsg,
    FullSyncStartMsg,
    FullSyncStartRetMsg,
    FullSyncStatusMsg,
    FullSyncStatusRetMsg,
    HealthWorkerMsg,
    HealthWorkerRetMsg,
    HeartbeatMsg,
    HeartbeatRetMsg,
    MoveWorkerMsg,
    MoveWorkerRetMsg,
    Msg,
    PinWorkerMsg,
    PinWorkerRetMsg,
    RegisterMsg,
    RegisterRetMsg,
    WorkerMsg,
    WorkerReqMsg,
    WorkerReqRetMsg,
)
from lmcache.v1.config import LMCacheEngineConfig
from lmcache.v1.metadata import LMCacheMetadata
from lmcache.v1.rpc_utils import (
    DEFAULT_SOCKET_RECV_TIMEOUT_MS,
    DEFAULT_SOCKET_SEND_TIMEOUT_MS,
    close_zmq_socket,
    get_ip,
    get_zmq_context,
    get_zmq_socket,
    get_zmq_socket_with_timeout,
)

if TYPE_CHECKING:
    # First Party
    from lmcache.v1.cache_engine import LMCacheEngine

logger = init_logger(__name__)


class LMCacheWorker:
    """
    LMCache Worker class to handle the execution of cache operations.
    This class is responsible for receiving requests from the executor and
    executing the corresponding operations on the LMCache engine.
    Each worker is associated with a specific LMCache instance and a worker id.
    """

    def __init__(
        self,
        config: LMCacheEngineConfig,
        metadata: LMCacheMetadata,
        lmcache_engine: "LMCacheEngine",
    ):
        # TODO (Jiayi): "instance_id" might not be needed anymore.
        # Please consider removing it.
        self.config = config
        self.lmcache_instance_id = config.lmcache_instance_id
        assert self.lmcache_instance_id is not None
        self.lmcache_engine = lmcache_engine
        self.worker_id = metadata.worker_id

        self.context = get_zmq_context()

        # Load timeout configurations from extra_config (in milliseconds)
        self.socket_recv_timeout_ms = config.get_extra_config_value(
            "worker_socket_recv_timeout_ms", DEFAULT_SOCKET_RECV_TIMEOUT_MS
        )
        self.socket_send_timeout_ms = config.get_extra_config_value(
            "worker_socket_send_timeout_ms", DEFAULT_SOCKET_SEND_TIMEOUT_MS
        )

        assert config.controller_pull_url is not None

        controller_pull_url = config.controller_pull_url
        self.push_socket = get_zmq_socket(
            self.context,
            controller_pull_url,
            protocol="tcp",
            role=zmq.PUSH,  # type: ignore[attr-defined]
            bind_or_connect="connect",
        )

        if config.controller_reply_url is not None:
            self.controller_rep_url = config.controller_reply_url
            self._create_req_socket()

        # Heartbeat socket will be created dynamically after register
        # based on heartbeat_url returned from controller
        self.heartbeat_socket: Optional[zmq.asyncio.Socket] = None
        self.controller_heartbeat_url: Optional[str] = None

        lmcache_worker_ids = config.get_lmcache_worker_ids(
            metadata.use_mla, metadata.world_size
        )
        if not lmcache_worker_ids:
            # start lmcache worker on all ranks
            assert len(config.lmcache_worker_ports) == metadata.world_size
            lmcache_worker_port = config.lmcache_worker_ports[self.worker_id]
        else:
            # start lmcache worker on given worker ids
            assert len(lmcache_worker_ids) == len(config.lmcache_worker_ports)
            index = lmcache_worker_ids.index(self.worker_id)
            lmcache_worker_port = config.lmcache_worker_ports[index]

        self.lmcache_worker_internal_url = f"*:{lmcache_worker_port}"
        self.lmcache_worker_ip = get_ip()
        self.lmcache_worker_port = lmcache_worker_port

        self.p2p_init_url = None
        if config.enable_p2p:
            self.p2p_host = config.p2p_host
            self.p2p_init_port = config.p2p_init_ports[self.worker_id]
            self.p2p_init_url = f"{self.p2p_host}:{self.p2p_init_port}"

        self.reply_socket = get_zmq_socket(
            self.context,
            self.lmcache_worker_internal_url,
            protocol="tcp",
            role=zmq.REP,  # type: ignore[attr-defined]
            bind_or_connect="bind",
        )

        logger.info(f"Reply socket established at {self.lmcache_worker_internal_url}")

        self.loop = asyncio.new_event_loop()
        self.thread = threading.Thread(
            target=self.loop.run_forever, daemon=True, name="lmcache-worker-thread"
        )
        self.thread.start()
        asyncio.run_coroutine_threadsafe(self.start_all(), self.loop)

        self.msg_queue: asyncio.Queue[WorkerMsg] = asyncio.Queue()

        # Full sync sender (initialized lazily when needed)
        self._full_sync_sender: Optional["FullSyncSender"] = None

    async def register(self):
        """
        Register the lmcache worker with the controller via DEALER-ROUTER.

        This method sends a RegisterMsg and waits for RegisterRetMsg
        which contains extra_config (e.g., heartbeat_url).
        """
        assert self.lmcache_instance_id is not None
        logger.info(
            "Registering lmcache instance-worker: %s",
            (self.lmcache_instance_id, self.worker_id),
        )

        register_msg = RegisterMsg(
            instance_id=self.lmcache_instance_id,
            worker_id=self.worker_id,
            ip=self.lmcache_worker_ip,
            port=self.lmcache_worker_port,
            peer_init_url=self.p2p_init_url,
        )

        # Send via DEALER socket (empty frame + payload) and wait for response
        try:
            await self.req_socket.send_multipart(
                [b"", msgspec.msgpack.encode(register_msg)]
            )
            # DEALER receives: [empty_frame, payload]
            frames = await self.req_socket.recv_multipart()
            serialized_ret_msg = frames[-1]
            ret_msg = msgspec.msgpack.decode(serialized_ret_msg, type=Msg)

            if isinstance(ret_msg, RegisterRetMsg):
                self._process_register_response(ret_msg)
            else:
                logger.warning("Unexpected register response type: %s", type(ret_msg))
        except zmq.ZMQError as e:
            logger.error("Failed to register with controller: %s", e)
            raise

    def _process_register_response(self, ret_msg: RegisterRetMsg):
        """Process RegisterRetMsg and initialize components based on extra_config."""
        extra_config = ret_msg.extra_config

        # Initialize heartbeat socket if heartbeat_url is provided
        heartbeat_url = extra_config.get("heartbeat_url")
        if heartbeat_url:
            logger.info("Received heartbeat_url from controller: %s", heartbeat_url)
            self.controller_heartbeat_url = heartbeat_url
            self._create_heartbeat_socket()
        else:
            logger.info("No dedicated heartbeat_url provided by controller")

    def deregister(self):
        """
        De-register the lmcache worker from the controller.
        """
        assert self.lmcache_instance_id is not None
        self.put_msg(
            DeRegisterMsg(
                instance_id=self.lmcache_instance_id,
                worker_id=self.worker_id,
                ip=self.lmcache_worker_ip,
                port=self.lmcache_worker_port,
            )
        )

    async def async_put_and_wait_msg(
        self,
        msg: WorkerReqMsg,
    ) -> WorkerReqRetMsg:
        """
        Send a message to the controller and wait for the response.

        This method handles different types of WorkerReqMsg using appropriate sockets:
        - HeartbeatMsg: Uses dedicated heartbeat socket
        - Other messages (RegisterMsg, BatchedP2PLookupMsg, FullSyncStartMsg,
          FullSyncStatusMsg): Uses DEALER socket (req_socket)

        Note: With DEALER-ROUTER mode, we no longer need to separate FullSync
        messages to heartbeat socket since DEALER supports async concurrent requests.
        """
        # Send heartbeat via dedicated heartbeat socket
        if isinstance(msg, HeartbeatMsg):
            return await self._send_heartbeat_msg(msg)

        # Send other messages via DEALER socket
        try:
            # DEALER socket: send [empty_frame, payload]
            await self.req_socket.send_multipart([b"", msgspec.msgpack.encode(msg)])
            frames = await self.req_socket.recv_multipart()
            # DEALER receives: [empty_frame, payload]
            serialized_ret_msg = frames[-1]
            ret_msg = msgspec.msgpack.decode(serialized_ret_msg, type=Msg)
            return ret_msg
        except zmq.Again as e:
            logger.error("Timeout occurred, recreating socket. Error: %s", e)
            self._recreate_req_socket()
            return self._on_request_failure(msg)
        except zmq.ZMQError as e:
            logger.error("ZMQ error occurred, recreating socket. Error: %s", e)
            self._recreate_req_socket()
            return self._on_request_failure(msg)
        except Exception as e:
            logger.error("Error happens in lmcache worker req_socket. Error: %s", e)
            return self._on_request_failure(msg)

    def _create_req_socket(self):
        self.req_socket = get_zmq_socket_with_timeout(
            self.context,
            self.controller_rep_url,
            "tcp",
            zmq.DEALER,  # type: ignore[attr-defined]
            "connect",
            self.socket_recv_timeout_ms,
            self.socket_send_timeout_ms,
        )

    def _create_heartbeat_socket(self):
        logger.info(
            "Creating heartbeat socket to connect to: %s, "
            "recv_timeout: %dms, send_timeout: %dms",
            self.controller_heartbeat_url,
            self.socket_recv_timeout_ms,
            self.socket_send_timeout_ms,
        )
        self.heartbeat_socket = get_zmq_socket_with_timeout(
            self.context,
            self.controller_heartbeat_url,
            "tcp",
            zmq.DEALER,  # type: ignore[attr-defined]
            "connect",
            self.socket_recv_timeout_ms,
            self.socket_send_timeout_ms,
        )

    def _recreate_heartbeat_socket(self):
        try:
            self.heartbeat_socket.close(linger=0)
        except Exception as e:
            logger.error("Error closing heartbeat socket: %s", e)
        self._create_heartbeat_socket()

    def _recreate_req_socket(self):
        try:
            self.req_socket.close(linger=0)
        except Exception as e:
            logger.error("Error closing req socket: %s", e)
        self._create_req_socket()

    def _get_full_sync_sender(self):
        """Lazy initialization of FullSyncSender"""
        if self._full_sync_sender is None:
            # Get the local_cpu_backend from lmcache_engine
            local_cpu_backend = self.lmcache_engine.storage_manager.local_cpu_backend
            self._full_sync_sender = FullSyncSender(
                config=self.config,
                worker=self,
                lmcache_engine=self.lmcache_engine,
                local_cpu_backend=local_cpu_backend,
            )
        return self._full_sync_sender

    def _on_request_failure(self, msg: WorkerReqMsg) -> WorkerReqRetMsg:
        """
        Create a default return message when worker -> controller
        request encounters an error (e.g., timeout, ZMQ error).
        """
        if isinstance(msg, BatchedP2PLookupMsg):
            return BatchedP2PLookupRetMsg(layout_info=[("", "", 0, "")])
        elif isinstance(msg, HeartbeatMsg):
            return HeartbeatRetMsg()  # No command by default
        elif isinstance(msg, FullSyncStartMsg):
            return FullSyncStartRetMsg(
                sync_id=msg.sync_id,
                accepted=False,
                error_msg="Communication error",
            )
        elif isinstance(msg, FullSyncStatusMsg):
            return FullSyncStatusRetMsg(
                sync_id=msg.sync_id,
                is_complete=False,
                global_progress=0.0,
                can_exit_freeze=False,
            )
        else:
            raise ValueError(f"Unknown message type: {type(msg)}")

    def put_msg(self, msg: WorkerMsg):
        """
        Put a message into the message queue.
        """
        # TODO(Jiayi): This might introduce ~0.05ms latency than
        # a normal function call.
        # Not sure how much overhead is blocking though.
        self.loop.call_soon_threadsafe(self.msg_queue.put_nowait, msg)

    async def batched_get_msg(self, max_bsz: int = 50) -> list[WorkerMsg]:
        """
        Get a batch of messages from the message queue.
        """
        batch = []

        # use blocking get for the first msg
        try:
            item = await self.msg_queue.get()
            batch.append(item)
        except asyncio.CancelledError:
            return batch  # shutdown path

        for _ in range(max_bsz - 1):
            try:
                item = self.msg_queue.get_nowait()
                batch.append(item)
            except asyncio.QueueEmpty:
                break
        return batch

    async def _send_heartbeat_msg(self, msg: HeartbeatMsg) -> HeartbeatRetMsg:
        """
        Send heartbeat message via dedicated heartbeat DEALER socket.
        This is separate from async_put_and_wait_msg to keep heartbeat independent.
        """
        if self.heartbeat_socket is None:
            logger.warning("Heartbeat socket is not initialized")
            return HeartbeatRetMsg()
        try:
            # DEALER socket: send [empty_frame, payload]
            logger.info("Sending heartbeat message to controller...")
            await self.heartbeat_socket.send_multipart(
                [b"", msgspec.msgpack.encode(msg)]
            )
            logger.info("Heartbeat message sent, waiting for response...")
            frames = await self.heartbeat_socket.recv_multipart()
            logger.info("Received heartbeat response with %d frames", len(frames))
            # DEALER receives: [empty_frame, payload]
            serialized_ret_msg = frames[-1]
            ret_msg = msgspec.msgpack.decode(serialized_ret_msg, type=Msg)
            return ret_msg
        except zmq.Again as e:
            logger.error("Heartbeat timeout occurred, recreating socket. Error: %s", e)
            self._recreate_heartbeat_socket()
            return HeartbeatRetMsg()
        except zmq.ZMQError as e:
            logger.error(
                "Heartbeat ZMQ error occurred, recreating socket. Error: %s", e
            )
            self._recreate_heartbeat_socket()
            return HeartbeatRetMsg()
        except Exception as e:
            logger.error("Error happens in heartbeat socket. Error: %s", e)
            return HeartbeatRetMsg()

    async def heartbeat(self):
        """
        Send periodic heartbeats to the controller (DEALER-ROUTER mode).

        Process any commands received in the heartbeat response.
        Uses dedicated heartbeat socket to avoid blocking from other requests.
        """
        enable_heartbeat = (
            self.config.lmcache_worker_heartbeat_time is not None
            and self.config.lmcache_worker_heartbeat_time > 0
            and self.heartbeat_socket is not None
        )
        if enable_heartbeat:
            await asyncio.sleep(self.config.lmcache_worker_heartbeat_delay_time)
            logger.info(
                "Start heartbeat in %s : %s, delay time: %ss, heartbeat time: %ss",
                self.lmcache_instance_id,
                self.worker_id,
                self.config.lmcache_worker_heartbeat_delay_time,
                self.config.lmcache_worker_heartbeat_time,
            )
            while True:
                # Send heartbeat via dedicated heartbeat socket
                heartbeat_msg = HeartbeatMsg(
                    instance_id=self.lmcache_instance_id,
                    worker_id=self.worker_id,
                    ip=self.lmcache_worker_ip,
                    port=self.lmcache_worker_port,
                    peer_init_url=self.p2p_init_url,
                )

                try:
                    ret_msg = await self._send_heartbeat_msg(heartbeat_msg)

                    if isinstance(ret_msg, HeartbeatRetMsg):
                        self._handle_heartbeat_commands(ret_msg)
                    else:
                        logger.warning(
                            "Unexpected heartbeat response type: %s", type(ret_msg)
                        )
                except Exception as e:
                    logger.error("Error during heartbeat: %s", e)

                await asyncio.sleep(self.config.lmcache_worker_heartbeat_time)

    def _handle_heartbeat_commands(self, ret_msg: HeartbeatRetMsg):
        """
        Handle commands received in heartbeat response.

        Uses polymorphic dispatch - each command class implements its own
        execute() method. Commands are executed sequentially.
        """
        if not ret_msg.has_commands():
            return

        for command in ret_msg.commands:
            logger.info(
                "Executing heartbeat command: %s",
                command.describe(),
            )
            try:
                command.execute(self)
            except NotImplementedError:
                logger.warning(
                    "Command %s.execute() not implemented yet",
                    type(command).__name__,
                )
            except Exception as e:
                logger.error(
                    "Error executing command %s: %s",
                    type(command).__name__,
                    e,
                )

    async def push(self):
        while True:
            try:
                msgs = await self.batched_get_msg()
                logger.debug(f"Sending {len(msgs)} messages")
                self.push_socket.send_multipart(
                    [msgspec.msgpack.encode(msg) for msg in msgs]
                )

            except Exception as e:
                logger.error(f"Push error: {e}")

    async def handle_request(self):
        """
        Handle incoming requests (control msgs) from the controller.
        """
        while True:
            try:
                serialized_request = await self.reply_socket.recv()
                request = msgspec.msgpack.decode(serialized_request, type=Msg)
                logger.debug(f"Received message: {request}")
                if isinstance(request, MoveWorkerMsg):
                    tokens = request.tokens
                    old_position = request.old_position
                    new_position = request.new_position
                    do_copy = request.copy
                    worker_event_id = request.worker_event_id

                    # Intra node move
                    if new_position[0] == self.lmcache_worker_internal_url:
                        # TODO(Jiayi): currently we only support moving from
                        # local disk to local cpu.
                        assert old_position[1] == "LocalDiskBackend"
                        assert new_position[1] == "LocalCPUBackend"
                        assert do_copy

                        # TODO(Jiayi): We need to align prefetch and move.
                        logger.debug("Executing prefetch operation.")
                        raise NotImplementedError(
                            "Prefetch from controller is not implemented yet."
                        )
                    else:
                        assert new_position[1] == "LocalCPUBackend", (
                            "Only support moving to cpu for now."
                        )
                        logger.debug("Executing cross-node move operation.")
                        num_tokens = self.lmcache_engine.move(
                            tokens=tokens,
                            old_position=old_position,
                            new_position=new_position,
                            event_id=worker_event_id,
                            do_copy=do_copy,
                        )

                    # TODO(Jiayi): LMCache needs to have an event tracking
                    # pool to enable more advanced control-plane optims.
                    # For now, we use a dummy `event_id`.
                    serialized_ret_msg = msgspec.msgpack.encode(
                        MoveWorkerRetMsg(num_tokens=num_tokens)
                    )
                elif isinstance(request, CompressWorkerMsg):
                    num_compressed_tokens = self.lmcache_engine.compress(
                        tokens=request.tokens,
                        method=request.method,
                        location=request.location,
                        event_id=request.worker_event_id,
                    )
                    serialized_ret_msg = msgspec.msgpack.encode(
                        CompressWorkerRetMsg(num_tokens=num_compressed_tokens)
                    )
                elif isinstance(request, DecompressWorkerMsg):
                    num_decompressed_tokens = self.lmcache_engine.decompress(
                        tokens=request.tokens,
                        method=request.method,
                        location=request.location,
                        event_id=request.worker_event_id,
                    )
                    serialized_ret_msg = msgspec.msgpack.encode(
                        DecompressWorkerRetMsg(num_tokens=num_decompressed_tokens)
                    )
                elif isinstance(request, PinWorkerMsg):
                    num_pinned_tokens = self.lmcache_engine.lookup(
                        tokens=request.tokens,
                        search_range=[request.location],
                        lookup_id=request.worker_event_id,
                        pin=True,
                    )
                    serialized_ret_msg = msgspec.msgpack.encode(
                        PinWorkerRetMsg(num_tokens=num_pinned_tokens)
                    )
                elif isinstance(request, ClearWorkerMsg):
                    num_cleared_tokens = self.lmcache_engine.clear(
                        locations=[request.location],
                    )
                    serialized_ret_msg = msgspec.msgpack.encode(
                        ClearWorkerRetMsg(num_tokens=num_cleared_tokens)
                    )
                elif isinstance(request, HealthWorkerMsg):
                    error_code = self.lmcache_engine.health()
                    serialized_ret_msg = msgspec.msgpack.encode(
                        HealthWorkerRetMsg(error_code=error_code)
                    )
                else:
                    logger.error(f"Unknown message: {request}")
                    serialized_ret_msg = msgspec.msgpack.encode(
                        ErrorMsg(error=f"Unknown message: {request}")
                    )

                await self.reply_socket.send(serialized_ret_msg)
            except Exception as e:
                logger.error(f"Worker error: {e}")
                serialized_ret_msg = msgspec.msgpack.encode(
                    ErrorMsg(error=f"Worker error: {e}")
                )
                await self.reply_socket.send(serialized_ret_msg)

    async def start_all(self):
        try:
            # Register first to get heartbeat_url before starting heartbeat task
            await self.register()

            logger.info(
                f"Starting lmcache worker {self.worker_id}"
                f"for instance {self.lmcache_instance_id}"
            )
            await asyncio.gather(
                self.push(),
                self.handle_request(),
                self.heartbeat(),
            )
        except Exception as e:
            logger.error(
                f"Instance {self.lmcache_instance_id}, "
                f"worker {self.worker_id} error: {e}"
            )

    def close(self):
        self.deregister()
        if self.loop.is_running():
            self.loop.call_soon_threadsafe(self.loop.stop)
        if self.thread.is_alive():
            self.thread.join()
        close_zmq_socket(self.push_socket)
        close_zmq_socket(self.reply_socket)
