"""Result types for message broker operations.

Analogous to ``db_core.query.queryset`` — ``MessageResult`` wraps a publish
acknowledgment and ``Subscription`` provides an async iterator for consuming
messages.
"""

from __future__ import annotations

import asyncio
import json
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, Optional


# ---------------------------------------------------------------------------
# MessageResult — publish acknowledgment
# ---------------------------------------------------------------------------


@dataclass
class MessageResult:
    """Wrapper around a publish acknowledgment from the broker.

    Attributes:
        topic: The topic the message was published to.
        partition: The partition the broker assigned.
        offset: The offset within the partition.
        timestamp: Broker-assigned timestamp in milliseconds.
        success: Whether the publish was acknowledged.
    """

    topic: str
    partition: int
    offset: int
    timestamp: int
    success: bool = True

    def failed(self) -> bool:
        """Return ``True`` if the publish was not acknowledged."""
        return not self.success

    def __repr__(self) -> str:
        return (
            f"MessageResult(topic={self.topic!r}, partition={self.partition}, "
            f"offset={self.offset}, timestamp={self.timestamp}, "
            f"success={self.success})"
        )


# ---------------------------------------------------------------------------
# ReceivedMessage — a single consumed message
# ---------------------------------------------------------------------------


@dataclass
class ReceivedMessage:
    """A single message received from a broker subscription.

    Attributes:
        topic: Source topic.
        partition: Source partition.
        offset: Offset within the partition.
        timestamp: Broker timestamp in milliseconds.
        key: Message key (may be ``None``).
        payload: Auto-deserialised JSON payload as a dictionary.
        headers: Message headers.
    """

    topic: str
    partition: int
    offset: int
    timestamp: int
    key: Optional[str]
    payload: Dict[str, Any]
    headers: Dict[str, str]
    _ack_callback: Optional[Callable] = field(default=None, repr=False)
    _nack_callback: Optional[Callable] = field(default=None, repr=False)

    async def ack(self) -> None:
        """Acknowledge successful processing (commit offset)."""
        if self._ack_callback is not None:
            await self._ack_callback(self)

    async def nack(self) -> None:
        """Signal processing failure (seek back for retry)."""
        if self._nack_callback is not None:
            await self._nack_callback(self)

    def __repr__(self) -> str:
        payload_preview = json.dumps(self.payload)
        if len(payload_preview) > 80:
            payload_preview = payload_preview[:77] + "..."
        return (
            f"ReceivedMessage(topic={self.topic!r}, partition={self.partition}, "
            f"offset={self.offset}, key={self.key!r}, payload={payload_preview})"
        )


# ---------------------------------------------------------------------------
# Subscription — async iterator / context manager for consuming
# ---------------------------------------------------------------------------


class Subscription:
    """Async iterator and context manager for consuming messages.

    Returned by ``.subscribe()`` and ``.listen()``.  Supports:

    * Async iteration: ``async for msg in subscription: ...``
    * Context manager: ``async with subscription as sub: ...``
    * Manual control: ``await subscription.stop()``

    The underlying consume loop runs in a background asyncio task.
    """

    def __init__(self, consumer: Any, topic: str):
        self._consumer = consumer
        self._topic = topic
        self._queue: asyncio.Queue[Optional[ReceivedMessage]] = asyncio.Queue()
        self._task: Optional[asyncio.Task] = None
        self._callback: Optional[Callable] = None
        self._stopped = False

    # -- Async context manager -------------------------------------------

    async def __aenter__(self) -> "Subscription":
        self._start_consume_loop()
        return self

    async def __aexit__(self, *args: Any) -> None:
        await self.stop()

    # -- Async iterator --------------------------------------------------

    def __aiter__(self) -> "Subscription":
        self._start_consume_loop()
        return self

    async def __anext__(self) -> ReceivedMessage:
        if self._stopped:
            raise StopAsyncIteration
        msg = await self._queue.get()
        if msg is None:
            # Sentinel — consume loop has stopped
            raise StopAsyncIteration
        return msg

    # -- Manual control --------------------------------------------------

    async def stop(self) -> None:
        """Signal the consumer loop to stop and wait for it to finish."""
        if self._stopped:
            return
        self._stopped = True
        if self._task is not None and not self._task.done():
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
        # Put sentinel to unblock any waiting __anext__
        await self._queue.put(None)

    async def commit(self) -> None:
        """Manually commit offsets (when auto-commit is disabled)."""
        if hasattr(self._consumer, "commit"):
            await self._consumer.commit()

    async def seek_to_beginning(self) -> None:
        """Seek to the earliest offset for all assigned partitions."""
        if hasattr(self._consumer, "seek_to_beginning"):
            await self._consumer.seek_to_beginning()

    async def seek_to_end(self) -> None:
        """Seek to the latest offset for all assigned partitions."""
        if hasattr(self._consumer, "seek_to_end"):
            await self._consumer.seek_to_end()

    # -- Internal --------------------------------------------------------

    def _start_callback(self, callback: Callable) -> None:
        """Assign a callback for ``.listen()`` mode."""
        self._callback = callback
        self._start_consume_loop()

    def _start_consume_loop(self) -> None:
        """Launch the background message polling task if not already running."""
        if self._task is not None and not self._task.done():
            return
        self._task = asyncio.ensure_future(self._consume_loop())

    async def _consume_loop(self) -> None:
        """Background loop: poll for messages from the consumer."""
        try:
            while not self._stopped:
                try:
                    raw = await asyncio.wait_for(
                        self._consumer.getone(), timeout=0.5
                    )
                except asyncio.TimeoutError:
                    continue

                if raw is None:
                    continue

                msg = self._wrap_message(raw)

                if self._callback is not None:
                    # Callback mode — invoke callback directly
                    try:
                        await self._callback(msg)
                    except Exception:
                        # Don't crash the consume loop on callback errors
                        pass
                else:
                    # Iterator mode — enqueue for __anext__
                    await self._queue.put(msg)
        except asyncio.CancelledError:
            pass
        except Exception:
            pass
        finally:
            if not self._stopped:
                await self._queue.put(None)  # sentinel for __anext__

    def _wrap_message(self, raw: Any) -> ReceivedMessage:
        """Wrap a raw consumer record into a :class:`ReceivedMessage`."""
        # raw is an aiokafka ConsumerRecord (or equivalent for other backends)

        # Extract payload — handle bytes → dict
        value = getattr(raw, "value", None)
        if isinstance(value, bytes):
            try:
                payload = json.loads(value.decode("utf-8"))
            except (json.JSONDecodeError, UnicodeDecodeError):
                payload = {"_raw": value.decode("utf-8", errors="replace")}
        elif isinstance(value, dict):
            payload = value
        else:
            payload = {"_value": str(value)} if value is not None else {}

        # Extract headers
        raw_headers = getattr(raw, "headers", ()) or ()
        headers = {}
        for h in raw_headers:
            if isinstance(h, tuple) and len(h) == 2:
                k, v = h
                headers[k.decode() if isinstance(k, bytes) else k] = (
                    v.decode() if isinstance(v, bytes) else v
                )

        # Extract key
        key = getattr(raw, "key", None)
        if isinstance(key, bytes):
            key = key.decode("utf-8")

        return ReceivedMessage(
            topic=getattr(raw, "topic", self._topic),
            partition=getattr(raw, "partition", -1),
            offset=getattr(raw, "offset", -1),
            timestamp=getattr(raw, "timestamp", 0),
            key=key,
            payload=payload,
            headers=headers,
            _ack_callback=self._make_ack_callback(raw),
            _nack_callback=self._make_nack_callback(raw),
        )

    def _make_ack_callback(self, raw: Any) -> Optional[Callable]:
        """Create an ack callback that commits the offset for this message."""

        async def _ack(msg: ReceivedMessage) -> None:
            if hasattr(self._consumer, "commit"):
                await self._consumer.commit()

        return _ack

    def _make_nack_callback(self, raw: Any) -> Optional[Callable]:
        """Create a nack callback (seek back for retry)."""

        async def _nack(msg: ReceivedMessage) -> None:
            if hasattr(self._consumer, "seek"):
                await self._consumer.seek(
                    partition=msg.partition, offset=msg.offset
                )

        return _nack
