from __future__ import annotations

from typing import Any

from .errors import MQConfigError
from .registry import MQRegistry
from .serialization import encode_message_key, encode_message_value


class MQClient:
    def __init__(self, producer: Any, registry: MQRegistry) -> None:
        self._producer = producer
        self._registry = registry

    def send_json(
        self,
        topic: str,
        value: Any,
        *,
        key: str | bytes | None = None,
        headers: dict[str, str | bytes | None] | None = None,
    ) -> None:
        self.send(topic, value=value, key=key, headers=headers, value_format="json")

    def send(
        self,
        topic: str,
        *,
        value: Any,
        key: str | bytes | None = None,
        headers: dict[str, str | bytes | None] | None = None,
        value_format: str = "bytes",
    ) -> None:
        self._producer.poll(0)
        self._producer.produce(
            topic,
            value=encode_message_value(value, value_format),
            key=encode_message_key(key),
            headers=_encode_headers(headers),
        )

    def sender(self, name: str) -> "MQSender":
        definition = self._registry.producers.get(name)
        if definition is None:
            raise MQConfigError(f"mq producer not registered: {name}")
        return MQSender(self, definition.topic, definition.value_format)

    def flush(self, timeout: float | None = None) -> None:
        if timeout is None:
            self._producer.flush()
        else:
            self._producer.flush(timeout)


class MQSender:
    def __init__(self, client: MQClient, topic: str, value_format: str) -> None:
        self._client = client
        self.topic = topic
        self.value_format = value_format

    def send_json(
        self,
        value: Any,
        *,
        key: str | bytes | None = None,
        headers: dict[str, str | bytes | None] | None = None,
    ) -> None:
        self._client.send(
            self.topic,
            value=value,
            key=key,
            headers=headers,
            value_format="json",
        )

    def send(
        self,
        value: Any,
        *,
        key: str | bytes | None = None,
        headers: dict[str, str | bytes | None] | None = None,
    ) -> None:
        self._client.send(
            self.topic,
            value=value,
            key=key,
            headers=headers,
            value_format=self.value_format,
        )


def _encode_headers(
    headers: dict[str, str | bytes | None] | None,
) -> list[tuple[str, bytes | None]] | None:
    if headers is None:
        return None
    result: list[tuple[str, bytes | None]] = []
    for key, value in headers.items():
        if value is None or isinstance(value, bytes):
            result.append((key, value))
        else:
            result.append((key, str(value).encode("utf-8")))
    return result


def mq_client(app) -> MQClient:
    client = getattr(app.state, "iti_mq_client", None)
    if client is None:
        raise MQConfigError("mq client is not configured")
    return client