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

# Third Party
import msgspec
import zmq

# First Party
from lmcache.logging import init_logger
from lmcache.v1.cache_controller.controllers import KVController, RegistrationController
from lmcache.v1.cache_controller.executor import LMCacheClusterExecutor
from lmcache.v1.cache_controller.observability import (
    PrometheusLogger,
    SocketMetricsContext,
    SocketType,
)
from lmcache.v1.rpc_utils import (
    get_ip,
    get_zmq_context,
    get_zmq_socket,
)

from lmcache.v1.cache_controller.message import (  # isort: skip
    BatchedKVOperationMsg,
    BatchedP2PLookupMsg,
    CheckFinishMsg,
    ClearMsg,
    CompressMsg,
    DecompressMsg,
    DeRegisterMsg,
    ErrorMsg,
    FullSyncBatchMsg,
    FullSyncEndMsg,
    FullSyncStartMsg,
    FullSyncStatusMsg,
    HealthMsg,
    HeartbeatMsg,
    LookupMsg,
    MoveMsg,
    Msg,
    MsgBase,
    OrchMsg,
    OrchRetMsg,
    PinMsg,
    QueryInstMsg,
    QueryWorkerInfoMsg,
    RegisterMsg,
    WorkerMsg,
    WorkerReqMsg,
    WorkerReqRetMsg,
)

logger = init_logger(__name__)

# TODO(Jiayi): Need to align the message types. For example,
# a controller should take in an control message and return
# a control message.


class LMCacheControllerManager:
    def __init__(
        self,
        controller_urls: dict[str, str],
        health_check_interval: int,
        lmcache_worker_timeout: int,
        full_sync_completion_threshold: float = 0.8,
        full_sync_timeout_s: float = 300.0,
    ):
        # Initialize stats logger
        prometheus_labels = {
            "role": "controller",
        }
        PrometheusLogger.GetOrCreate(prometheus_labels)
        self.zmq_context = get_zmq_context()
        self.controller_urls = controller_urls
        # TODO(Jiayi): We might need multiple sockets if there are more
        # controllers. For now, we use a single socket to receive messages
        # for all controllers.
        # Similarly we might need more sockets to handle different control
        # messages. For now, we use one socket to handle all control messages.

        # TODO(Jiayi): Another thing is that we might need to decoupe the
        # interactions among `handle_worker_message`, `handle_control_message`
        # and `handle_orchestration_message`. For example, in
        # `handle_orchestration_message`, we might need to call
        # `issue_control_message`. This will make the system less concurrent.

        # Micro controllers
        self.controller_pull_socket = get_zmq_socket(
            self.zmq_context,
            self.controller_urls["pull"],
            protocol="tcp",
            role=zmq.PULL,  # type: ignore[attr-defined]
            bind_or_connect="bind",
        )

        if self.controller_urls["reply"] is not None:
            self.controller_reply_socket = get_zmq_socket(
                self.zmq_context,
                self.controller_urls["reply"],
                protocol="tcp",
                role=zmq.ROUTER,  # type: ignore[attr-defined]
                bind_or_connect="bind",
            )

        # Dedicated heartbeat socket to avoid blocking from other requests
        if self.controller_urls.get("heartbeat") is not None:
            self.controller_heartbeat_socket = get_zmq_socket(
                self.zmq_context,
                self.controller_urls["heartbeat"],
                protocol="tcp",
                role=zmq.ROUTER,  # type: ignore[attr-defined]
                bind_or_connect="bind",
            )
        else:
            self.controller_heartbeat_socket = None
        self.reg_controller = RegistrationController()
        self.kv_controller = KVController(
            registry=self.reg_controller.registry,
            full_sync_completion_threshold=full_sync_completion_threshold,
            full_sync_timeout_s=full_sync_timeout_s,
        )

        # Cluster executor
        self.cluster_executor = LMCacheClusterExecutor(
            reg_controller=self.reg_controller,
        )

        # post initialization of controllers
        self.kv_controller.post_init(
            reg_controller=self.reg_controller,
            cluster_executor=self.cluster_executor,
        )
        self.reg_controller.post_init(
            kv_controller=self.kv_controller,
            cluster_executor=self.cluster_executor,
        )
        self.health_check_interval = health_check_interval
        self.lmcache_worker_timeout = lmcache_worker_timeout

        if self.health_check_interval > 0:
            logger.info(
                "Start health check thread, interval: %s", self.health_check_interval
            )
            self.loop = asyncio.new_event_loop()
            self.thread = threading.Thread(
                target=self.loop.run_forever,
                daemon=True,
                name="controller-health-thread",
            )
            self.thread.start()
            asyncio.run_coroutine_threadsafe(self.health_check(), self.loop)

        # Setup socket message count metrics
        self._setup_socket_metrics()

    async def handle_worker_message(self, msg: WorkerMsg) -> None:
        if isinstance(msg, RegisterMsg):
            await self.reg_controller.register(msg)
        elif isinstance(msg, DeRegisterMsg):
            await self.reg_controller.deregister(msg)
        elif isinstance(msg, BatchedKVOperationMsg):
            await self.kv_controller.handle_batched_kv_operations(msg)
        elif isinstance(msg, FullSyncBatchMsg):
            await self.kv_controller.handle_full_sync_batch(msg)
        elif isinstance(msg, FullSyncEndMsg):
            await self.kv_controller.handle_full_sync_end(msg)
        else:
            logger.error(f"Unknown worker message type: {msg}")

    async def handle_worker_req_message(
        self, msg: WorkerReqMsg
    ) -> Union[WorkerReqRetMsg, ErrorMsg]:
        ret_msg: Union[WorkerReqRetMsg, ErrorMsg]
        if isinstance(msg, RegisterMsg):
            # Build extra_config with heartbeat_url if available
            extra_config: dict[str, str] = {}
            if self.controller_urls.get("heartbeat") is not None:
                # Convert bind address (e.g., "0.0.0.0:8082" or "*:8082")
                # to a connectable address using actual controller IP
                heartbeat_bind_url = self.controller_urls["heartbeat"]
                heartbeat_url = self._convert_bind_to_connect_url(
                    heartbeat_bind_url, worker_ip=msg.ip
                )
                extra_config["heartbeat_url"] = heartbeat_url
                logger.debug(
                    "Returning heartbeat_url to worker: %s (bind: %s)",
                    heartbeat_url,
                    heartbeat_bind_url,
                )
            ret_msg = await self.reg_controller.register(msg, extra_config)
        elif isinstance(msg, BatchedP2PLookupMsg):
            ret_msg = await self.kv_controller.batched_p2p_lookup(msg)
        elif isinstance(msg, HeartbeatMsg):
            ret_msg = await self.reg_controller.heartbeat(msg)
        elif isinstance(msg, FullSyncStartMsg):
            ret_msg = await self.kv_controller.handle_full_sync_start(msg)
        elif isinstance(msg, FullSyncStatusMsg):
            ret_msg = await self.kv_controller.handle_full_sync_status(msg)
        else:
            logger.error(f"Unknown worker request message type: {msg}")
            ret_msg = ErrorMsg(error=f"Unknown message type: {type(msg)}")
        return ret_msg

    async def handle_orchestration_message(self, msg: OrchMsg) -> OrchRetMsg:
        if isinstance(msg, LookupMsg):
            return await self.kv_controller.lookup(msg)
        elif isinstance(msg, HealthMsg):
            return await self.reg_controller.health(msg)
        elif isinstance(msg, QueryInstMsg):
            return await self.reg_controller.get_instance_id(msg)
        elif isinstance(msg, ClearMsg):
            return await self.kv_controller.clear(msg)
        elif isinstance(msg, PinMsg):
            return await self.kv_controller.pin(msg)
        elif isinstance(msg, CompressMsg):
            return await self.kv_controller.compress(msg)
        elif isinstance(msg, DecompressMsg):
            return await self.kv_controller.decompress(msg)
        elif isinstance(msg, MoveMsg):
            return await self.kv_controller.move(msg)
        elif isinstance(msg, CheckFinishMsg):
            # FIXME(Jiayi): This `check_finish` thing
            # shouldn't be implemented in kv_controller.
            return await self.kv_controller.check_finish(msg)
        elif isinstance(msg, QueryWorkerInfoMsg):
            return await self.reg_controller.query_worker_info(msg)
        else:
            logger.error(f"Unknown orchestration message type: {msg}")
            raise RuntimeError(f"Unknown orchestration message type: {msg}")

    def _setup_socket_metrics(self):
        """Setup metrics for socket message counts."""
        # Initialize message counters for observability
        self.pull_socket_message_count = 0
        self.reply_socket_message_count = 0

        # Initialize active request counters
        self.pull_socket_active_requests = 0
        self.reply_socket_active_requests = 0

        prometheus_logger = PrometheusLogger.GetInstanceOrNone()
        if prometheus_logger is not None:
            prometheus_logger.pull_socket_message_count.set_function(
                lambda: self.pull_socket_message_count
            )
            prometheus_logger.reply_socket_message_count.set_function(
                lambda: self.reply_socket_message_count
            )

            # Socket queue/backlog metrics
            prometheus_logger.pull_socket_has_pending.set_function(
                lambda: self._check_socket_has_pending(self.controller_pull_socket)
            )
            if self.controller_urls["reply"] is not None:
                prometheus_logger.reply_socket_has_pending.set_function(
                    lambda: self._check_socket_has_pending(self.controller_reply_socket)
                )

            # Active request metrics
            prometheus_logger.pull_socket_active_requests.set_function(
                lambda: self.pull_socket_active_requests
            )
            prometheus_logger.reply_socket_active_requests.set_function(
                lambda: self.reply_socket_active_requests
            )

    def _convert_bind_to_connect_url(
        self, bind_url: str, worker_ip: Optional[str] = None
    ) -> str:
        """Convert a bind address to a connectable address.

        Bind addresses like "0.0.0.0:port" or "*:port" cannot be used
        by workers to connect. We need to replace them with the actual
        controller IP address.

        If worker_ip is provided and matches the controller's IP, use
        127.0.0.1 for loopback connection (more reliable than external IP).

        Args:
            bind_url: The bind URL (e.g., "0.0.0.0:8082" or "*:8082")
            worker_ip: The IP address of the requesting worker (optional)

        Returns:
            A connectable URL (e.g., "192.168.1.100:8082" or "127.0.0.1:8082")
        """
        if ":" not in bind_url:
            return bind_url

        host, port = bind_url.rsplit(":", 1)
        # Replace bind-all addresses with actual IP
        if host in ("0.0.0.0", "*", ""):
            actual_ip = get_ip()
            # If worker is on the same machine, use loopback for reliability
            # This handles cases where external IP (e.g., VPN) doesn't support
            # local loopback connections
            if worker_ip is not None and worker_ip == actual_ip:
                logger.debug(
                    "Worker IP %s matches controller IP, using 127.0.0.1",
                    worker_ip,
                )
                return f"127.0.0.1:{port}"
            return f"{actual_ip}:{port}"
        return bind_url

    def _check_socket_has_pending(self, socket) -> int:
        """Check if socket has pending messages.

        Returns:
            1 if socket has pending messages, 0 otherwise
        """
        try:
            events = socket.get(zmq.EVENTS)  # type: ignore[attr-defined]
            # Check if POLLIN flag is set (indicates readable/pending messages)
            has_pending = 1 if (events & zmq.POLLIN) else 0  # type: ignore[attr-defined]
            return has_pending
        except Exception as e:
            logger.error(f"Error checking socket pending status: {e}")
            return 0

    async def handle_batched_push_request(self, socket) -> Optional[MsgBase]:
        while True:
            parts = await socket.recv_multipart()
            part_count = len(parts)
            with SocketMetricsContext(self, SocketType.PULL, part_count):
                for part in parts:
                    # Parse message based on format
                    if part.startswith(b"{"):
                        # JSON format - typically from external systems
                        # like Mooncake
                        msg_dict = json.loads(part)
                        msg = msgspec.convert(msg_dict, type=Msg)
                    else:
                        # MessagePack format - internal LMCache communication
                        msg = msgspec.msgpack.decode(part, type=Msg)
                    if isinstance(msg, WorkerMsg):
                        await self.handle_worker_message(msg)

                    # FIXME(Jiayi): The abstraction of control messages
                    # might not be necessary.
                    # elif isinstance(msg, ControlMsg):
                    #    await self.issue_control_message(msg)
                    elif isinstance(msg, OrchMsg):
                        await self.handle_orchestration_message(msg)
                    else:
                        logger.error(f"Unknown message type: {type(msg)}")

    async def handle_batched_req_request(self, socket) -> Optional[MsgBase]:
        """Handle requests on ROUTER socket.

        ROUTER socket receives multi-part messages:
        [identity, empty_frame, payload]
        and must reply with the same identity frame.
        """
        while True:
            frames = await socket.recv_multipart()
            with SocketMetricsContext(self, SocketType.REPLY):
                identity = None
                try:
                    # ROUTER socket: [identity, empty_frame, payload]
                    if len(frames) < 3:
                        logger.error(
                            "Invalid ROUTER message format, expected >= 3 frames, "
                            "got %d",
                            len(frames),
                        )
                        continue
                    identity = frames[0]
                    # frames[1] is empty delimiter
                    part = frames[2]

                    # Parse message based on format
                    if part.startswith(b"{"):
                        # JSON format - typically from external systems like Mooncake
                        msg_dict = json.loads(part)
                        msg = msgspec.convert(msg_dict, type=Msg)
                    else:
                        # MessagePack format - internal LMCache communication
                        msg = msgspec.msgpack.decode(part, type=Msg)

                    if isinstance(msg, WorkerReqMsg):
                        ret_msg = await self.handle_worker_req_message(msg)
                        # Reply with identity frame for ROUTER socket
                        await socket.send_multipart(
                            [identity, b"", msgspec.msgpack.encode(ret_msg)]
                        )
                    else:
                        logger.error("Unknown message type: %s", type(msg))
                        err_msg = ErrorMsg(error=f"Unknown message type: {type(msg)}")
                        await socket.send_multipart(
                            [identity, b"", msgspec.msgpack.encode(err_msg)]
                        )
                except (
                    json.JSONDecodeError,
                    msgspec.DecodeError,
                    msgspec.ValidationError,
                    zmq.ZMQError,
                ) as e:
                    logger.error("Error handling request message: %s", e, exc_info=True)
                    err_msg = ErrorMsg(error=str(e))
                    # Try to reply with error if we have identity
                    if identity is not None:
                        await socket.send_multipart(
                            [identity, b"", msgspec.msgpack.encode(err_msg)]
                        )

    async def handle_heartbeat_request(self, socket) -> None:
        """Handle heartbeat requests on dedicated ROUTER socket.

        This runs on a separate socket to ensure heartbeats are processed
        without being blocked by other requests.

        ROUTER socket receives multi-part messages:
        [identity, empty_frame, payload]
        """
        logger.info("Heartbeat handler task started, waiting for heartbeat requests...")
        while True:
            frames = await socket.recv_multipart()
            logger.debug("Received heartbeat request with %d frames", len(frames))
            with SocketMetricsContext(self, SocketType.REPLY):
                identity = None
                try:
                    # ROUTER socket: [identity, empty_frame, payload]
                    if len(frames) < 3:
                        logger.error(
                            "Invalid heartbeat ROUTER message format, "
                            "expected >= 3 frames, got %d",
                            len(frames),
                        )
                        continue
                    identity = frames[0]
                    part = frames[2]

                    if part.startswith(b"{"):
                        msg_dict = json.loads(part)
                        msg = msgspec.convert(msg_dict, type=Msg)
                    else:
                        msg = msgspec.msgpack.decode(part, type=Msg)

                    if isinstance(msg, HeartbeatMsg):
                        ret_msg = await self.reg_controller.heartbeat(msg)
                        await socket.send_multipart(
                            [identity, b"", msgspec.msgpack.encode(ret_msg)]
                        )
                    else:
                        logger.error(
                            "Unexpected message type on heartbeat socket: %s",
                            type(msg),
                        )
                        err_msg = ErrorMsg(
                            error=f"Expected HeartbeatMsg, got {type(msg)}"
                        )
                        await socket.send_multipart(
                            [identity, b"", msgspec.msgpack.encode(err_msg)]
                        )
                except (
                    json.JSONDecodeError,
                    msgspec.DecodeError,
                    msgspec.ValidationError,
                    zmq.ZMQError,
                ) as e:
                    logger.error(
                        "Error handling heartbeat request: %s", e, exc_info=True
                    )
                    err_msg = ErrorMsg(error=str(e))
                    # Try to reply with error if we have identity
                    if identity is not None:
                        await socket.send_multipart(
                            [identity, b"", msgspec.msgpack.encode(err_msg)]
                        )

    async def health_check(self):
        while True:
            await asyncio.sleep(self.health_check_interval)
            worker_infos = self.reg_controller.registry.get_all_worker_infos_cached(1)
            for worker_info in worker_infos:
                if (
                    time.time() - worker_info.last_heartbeat_time
                    > self.lmcache_worker_timeout
                ):
                    logger.warning(
                        "Worker %s_%s last heartbeat time: %s, "
                        "current time: %s, more than %s seconds",
                        worker_info.instance_id,
                        worker_info.worker_id,
                        worker_info.last_heartbeat_time,
                        time.time(),
                        self.lmcache_worker_timeout,
                    )
                    # Perform a full deregister to clean up all associated resources.
                    deregister_msg = DeRegisterMsg(
                        instance_id=worker_info.instance_id,
                        worker_id=worker_info.worker_id,
                        ip=worker_info.ip,
                        port=worker_info.port,
                    )
                    await self.reg_controller.deregister(deregister_msg)

    async def start_all(self):
        tasks = []
        if self.controller_urls["reply"] is not None:
            logger.info(
                "Starting reply request handler on %s", self.controller_urls["reply"]
            )
            tasks.append(self.handle_batched_req_request(self.controller_reply_socket))
        if self.controller_heartbeat_socket is not None:
            logger.info(
                "Starting heartbeat request handler on %s",
                self.controller_urls.get("heartbeat"),
            )
            tasks.append(
                self.handle_heartbeat_request(self.controller_heartbeat_socket)
            )
        else:
            logger.warning("Heartbeat socket is None, heartbeat handler not started!")
        logger.info("Starting pull request handler on %s", self.controller_urls["pull"])
        tasks.append(self.handle_batched_push_request(self.controller_pull_socket))
        await asyncio.gather(
            *tasks,
            return_exceptions=True,
        )
