# Copyright 2025-2026 Dimensional Inc.
#
# 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.

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from dimos.protocol.pubsub.spec import PickleEncoderMixin, PubSub, PubSubEncoderMixin
from dimos.protocol.service.lcmservice import (
    LCMConfig,
    LCMService,
    autoconf,
)
from dimos.utils.logging_config import setup_logger

if TYPE_CHECKING:
    from collections.abc import Callable
    import threading

    from dimos.msgs import DimosMsg

logger = setup_logger()


@dataclass
class Topic:
    topic: str = ""
    lcm_type: type[DimosMsg] | None = None

    def __str__(self) -> str:
        if self.lcm_type is None:
            return self.topic
        return f"{self.topic}#{self.lcm_type.msg_name}"


class LCMPubSubBase(LCMService, PubSub[Topic, Any]):
    default_config = LCMConfig
    _stop_event: threading.Event
    _thread: threading.Thread | None
    _callbacks: dict[str, list[Callable[[Any], None]]]

    def __init__(self, **kwargs) -> None:  # type: ignore[no-untyped-def]
        super().__init__(**kwargs)
        self._callbacks = {}

    def publish(self, topic: Topic, message: bytes) -> None:
        """Publish a message to the specified channel."""
        if self.l is None:
            logger.error("Tried to publish after LCM was closed")
            return

        self.l.publish(str(topic), message)

    def subscribe(
        self, topic: Topic, callback: Callable[[bytes, Topic], Any]
    ) -> Callable[[], None]:
        if self.l is None:
            logger.error("Tried to subscribe after LCM was closed")

            def noop() -> None:
                pass

            return noop

        lcm_subscription = self.l.subscribe(str(topic), lambda _, msg: callback(msg, topic))

        # Set queue capacity to 10000 to handle high-volume bursts
        lcm_subscription.set_queue_capacity(10000)

        def unsubscribe() -> None:
            if self.l is None:
                return
            self.l.unsubscribe(lcm_subscription)

        return unsubscribe


class LCMEncoderMixin(PubSubEncoderMixin[Topic, Any, bytes]):
    def encode(self, msg: DimosMsg, _: Topic) -> bytes:
        return msg.lcm_encode()

    def decode(self, msg: bytes, topic: Topic) -> DimosMsg:
        if topic.lcm_type is None:
            raise ValueError(
                f"Cannot decode message for topic '{topic.topic}': no lcm_type specified"
            )
        return topic.lcm_type.lcm_decode(msg)


class JpegEncoderMixin(PubSubEncoderMixin[Topic, Any, bytes]):
    def encode(self, msg: DimosMsg, _: Topic) -> bytes:
        return msg.lcm_jpeg_encode()  # type: ignore[attr-defined, no-any-return]

    def decode(self, msg: bytes, topic: Topic) -> DimosMsg:
        if topic.lcm_type is None:
            raise ValueError(
                f"Cannot decode message for topic '{topic.topic}': no lcm_type specified"
            )
        return topic.lcm_type.lcm_jpeg_decode(msg)  # type: ignore[attr-defined, no-any-return]


class LCM(
    LCMEncoderMixin,
    LCMPubSubBase,
): ...


class PickleLCM(
    PickleEncoderMixin,  # type: ignore[type-arg]
    LCMPubSubBase,
): ...


class JpegLCM(
    JpegEncoderMixin,
    LCMPubSubBase,
): ...


__all__ = [
    "LCM",
    "JpegLCM",
    "LCMEncoderMixin",
    "LCMPubSubBase",
    "PickleLCM",
    "autoconf",
]
