# config/rabbitmq.py
# One long-lived connection per process, multiple channels on top.
# We use aio-pika because it's the idiomatic async AMQP client for Python --
# pika (sync) works too, but blocks your event loop if you're running FastAPI
# or any asyncio-based service.
import asyncio
import logging
import os
from typing import Optional

import aio_pika
from aio_pika import ExchangeType
from aio_pika.abc import AbstractRobustConnection, AbstractRobustChannel

logger = logging.getLogger(__name__)


class RabbitMQConnection:
    def __init__(self) -> None:
        self.connection: Optional[AbstractRobustConnection] = None
        self.channel: Optional[AbstractRobustChannel] = None

    async def connect(self) -> AbstractRobustChannel:
        try:
            # connect_robust auto-reconnects on network errors --
            # this replaces the manual reconnect loop you'd write with pika.
            self.connection = await aio_pika.connect_robust(
                host=os.getenv("RABBITMQ_HOST", "localhost"),
                port=int(os.getenv("RABBITMQ_PORT", "5672")),
                login=os.getenv("RABBITMQ_USER", "guest"),
                password=os.getenv("RABBITMQ_PASS", "guest"),
                virtualhost=os.getenv("RABBITMQ_VHOST", "/"),
                heartbeat=60,  # detect half-open connections
            )
            self.channel = await self.connection.channel()
            # publisher confirms: the broker ACKs every publish.
            # Without this, a publish can silently fail on broker crash.
            await self.channel.set_qos(prefetch_count=10)

            self.connection.close_callbacks.add(self._on_close)
            logger.info("Connected to RabbitMQ")
            return self.channel
        except Exception:
            logger.exception("Failed to connect to RabbitMQ")
            raise

    def _on_close(self, connection, exc) -> None:
        logger.warning("RabbitMQ connection closed: %s", exc)

    async def setup_exchanges(self) -> None:
        assert self.channel is not None
        # Durable=True means the exchange survives broker restarts.
        await self.channel.declare_exchange(
            "direct_exchange", ExchangeType.DIRECT, durable=True
        )
        await self.channel.declare_exchange(
            "topic_exchange", ExchangeType.TOPIC, durable=True
        )
        await self.channel.declare_exchange(
            "fanout_exchange", ExchangeType.FANOUT, durable=True
        )
        # Dead-letter exchange: monitor this; messages here = humans needed.
        await self.channel.declare_exchange(
            "dlx_exchange", ExchangeType.DIRECT, durable=True
        )

    async def close(self) -> None:
        if self.channel:
            await self.channel.close()
        if self.connection:
            await self.connection.close()


rabbitmq = RabbitMQConnection()

# messaging/producer.py
import json
import os
import uuid
from datetime import datetime, timezone
from typing import Any, Optional

import aio_pika
from aio_pika import DeliveryMode, Message
from aio_pika.abc import AbstractRobustChannel
from pydantic import BaseModel, Field


class DomainEvent(BaseModel):
    """Canonical event envelope. Use Pydantic so every producer emits the
    same shape -- consumers rely on this contract."""
    event_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    event_type: str
    occurred_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    version: str = "1.0"
    correlation_id: Optional[str] = None
    data: dict[str, Any]
    metadata: dict[str, Any] = Field(default_factory=dict)


class MessageProducer:
    def __init__(self, channel: AbstractRobustChannel) -> None:
        self.channel = channel
        self.service_name = os.getenv("SERVICE_NAME", "unknown")

    async def send_to_queue(
        self,
        queue_name: str,
        payload: dict[str, Any],
        headers: Optional[dict[str, Any]] = None,
    ) -> None:
        # Declare queue with a DLX binding so poison messages land in DLQ.
        queue = await self.channel.declare_queue(
            queue_name,
            durable=True,
            arguments={
                "x-dead-letter-exchange": "dlx_exchange",
                "x-dead-letter-routing-key": f"{queue_name}.dlq",
            },
        )

        msg = Message(
            body=json.dumps(payload).encode(),
            delivery_mode=DeliveryMode.PERSISTENT,  # survives broker restart
            message_id=payload.get("id") or str(uuid.uuid4()),
            timestamp=datetime.now(timezone.utc),
            content_type="application/json",
            headers={
                "x-retry-count": 0,
                "x-source-service": self.service_name,
                **(headers or {}),
            },
        )
        # default exchange ("") routes by queue name -- this is point-to-point.
        await self.channel.default_exchange.publish(msg, routing_key=queue.name)

    async def publish(
        self,
        exchange_name: str,
        routing_key: str,
        event: DomainEvent,
        extra_headers: Optional[dict[str, Any]] = None,
    ) -> None:
        exchange = await self.channel.get_exchange(exchange_name)
        msg = Message(
            body=event.model_dump_json().encode(),
            delivery_mode=DeliveryMode.PERSISTENT,
            message_id=event.event_id,
            timestamp=event.occurred_at,
            content_type="application/json",
            type=event.event_type,
            headers={
                "x-event-version": event.version,
                "x-correlation-id": event.correlation_id or "",
                "x-source-service": self.service_name,
                **(extra_headers or {}),
            },
        )
        await exchange.publish(msg, routing_key=routing_key)

    async def publish_domain_event(self, event: DomainEvent) -> None:
        # Routing key = event type, e.g., "order.created" -> topic subscribers
        # bound on "order.*" or "order.#" receive it.
        event.metadata.setdefault("service", self.service_name)
        event.metadata.setdefault("environment", os.getenv("ENV", "dev"))
        await self.publish("topic_exchange", event.event_type, event)

# messaging/consumer.py
import asyncio
import json
import logging
import os
from datetime import datetime, timezone
from typing import Any, Awaitable, Callable, Iterable

import aio_pika
from aio_pika.abc import AbstractIncomingMessage, AbstractRobustChannel

logger = logging.getLogger(__name__)

Handler = Callable[[dict[str, Any], AbstractIncomingMessage], Awaitable[None]]


class MessageConsumer:
    def __init__(self, channel: AbstractRobustChannel) -> None:
        self.channel = channel
        self.service_name = os.getenv("SERVICE_NAME", "unknown")

    async def consume_queue(
        self,
        queue_name: str,
        handler: Handler,
        prefetch: int = 10,
        max_retries: int = 3,
    ) -> None:
        # Prefetch caps the in-flight messages per consumer -- crucial for
        # fair dispatch when consumers process at different speeds.
        await self.channel.set_qos(prefetch_count=prefetch)

        queue = await self.channel.declare_queue(
            queue_name,
            durable=True,
            arguments={
                "x-dead-letter-exchange": "dlx_exchange",
                "x-dead-letter-routing-key": f"{queue_name}.dlq",
            },
        )
        logger.info("Consuming from queue: %s", queue_name)

        async def _on_message(message: AbstractIncomingMessage) -> None:
            started = datetime.now(timezone.utc)
            message_id = message.message_id or "<no-id>"
            # async with message.process() would auto-ack; we handle acks manually
            # because we need retry logic.
            try:
                payload = json.loads(message.body)
                await handler(payload, message)
                await message.ack()
                elapsed = (datetime.now(timezone.utc) - started).total_seconds() * 1000
                logger.info("Processed %s in %.1fms", message_id, elapsed)
            except Exception:
                logger.exception("Error processing message %s", message_id)
                retry_count = int((message.headers or {}).get("x-retry-count", 0)) + 1
                if retry_count <= max_retries:
                    await self._retry(queue_name, message, retry_count)
                    await message.ack()
                else:
                    # nack with requeue=False routes to the DLX.
                    await message.nack(requeue=False)
                    logger.warning(
                        "Message %s sent to DLQ after %d retries",
                        message_id,
                        max_retries,
                    )

        await queue.consume(_on_message)

    async def subscribe_to_topic(
        self,
        patterns: Iterable[str],
        handler: Handler,
        queue_name: str | None = None,
        prefetch: int = 10,
    ) -> None:
        patterns = list(patterns)
        qname = queue_name or (
            f"{self.service_name}."
            + "_".join(p.replace(".", "_").replace("#", "_").replace("*", "_") for p in patterns)
        )
        queue = await self.channel.declare_queue(
            qname,
            durable=True,
            arguments={
                "x-dead-letter-exchange": "dlx_exchange",
                "x-dead-letter-routing-key": f"{qname}.dlq",
            },
        )
        exchange = await self.channel.get_exchange("topic_exchange")
        for pattern in patterns:
            await queue.bind(exchange, routing_key=pattern)
            logger.info("Bound %s to pattern %s", qname, pattern)

        await self.consume_queue(qname, handler, prefetch=prefetch)

    async def _retry(
        self,
        queue_name: str,
        message: AbstractIncomingMessage,
        retry_count: int,
    ) -> None:
        # Exponential backoff: 2s, 4s, 8s, ...
        delay = 2**retry_count
        await asyncio.sleep(delay)
        new_headers = dict(message.headers or {})
        new_headers.update(
            {
                "x-retry-count": retry_count,
                "x-retry-timestamp": datetime.now(timezone.utc).isoformat(),
            }
        )
        await self.channel.default_exchange.publish(
            aio_pika.Message(
                body=message.body,
                headers=new_headers,
                content_type=message.content_type,
                message_id=message.message_id,
                delivery_mode=aio_pika.DeliveryMode.PERSISTENT,
            ),
            routing_key=queue_name,
        )

# order_service/events.py
from decimal import Decimal
from typing import Any

from messaging.producer import DomainEvent, MessageProducer


class OrderEventPublisher:
    def __init__(self, producer: MessageProducer) -> None:
        self.producer = producer

    async def order_created(self, order: dict[str, Any]) -> None:
        await self.producer.publish_domain_event(
            DomainEvent(
                event_type="order.created",
                correlation_id=order.get("correlation_id"),
                data={
                    "order_id": order["id"],
                    "customer_id": order["customer_id"],
                    "items": [
                        {
                            "product_id": item["product_id"],
                            "quantity": item["quantity"],
                            "price": str(item["price"]),  # Decimal -> str for JSON
                        }
                        for item in order["items"]
                    ],
                    "total_amount": str(order["total_amount"]),
                    "shipping_address": order["shipping_address"],
                    "created_at": order["created_at"].isoformat(),
                },
            )
        )

    async def order_paid(self, order_id: str, payment_id: str, amount: Decimal) -> None:
        await self.producer.publish_domain_event(
            DomainEvent(
                event_type="order.paid",
                data={"order_id": order_id, "payment_id": payment_id, "amount": str(amount)},
            )
        )

    async def order_shipped(
        self, order_id: str, tracking_number: str, carrier: str
    ) -> None:
        await self.producer.publish_domain_event(
            DomainEvent(
                event_type="order.shipped",
                data={
                    "order_id": order_id,
                    "tracking_number": tracking_number,
                    "carrier": carrier,
                },
            )
        )

    async def order_cancelled(self, order_id: str, reason: str) -> None:
        await self.producer.publish_domain_event(
            DomainEvent(
                event_type="order.cancelled",
                data={"order_id": order_id, "reason": reason},
            )
        )


# inventory_service/event_handlers.py
from aio_pika.abc import AbstractIncomingMessage

from messaging.consumer import MessageConsumer


class InventoryEventHandler:
    def __init__(self, consumer: MessageConsumer, inventory_service) -> None:
        self.consumer = consumer
        self.inventory = inventory_service

    async def start(self) -> None:
        await self.consumer.subscribe_to_topic(
            patterns=["order.created", "order.cancelled"],
            handler=self.handle_order_event,
            queue_name="inventory-service.orders",
        )

    async def handle_order_event(
        self, event: dict[str, Any], msg: AbstractIncomingMessage
    ) -> None:
        event_type = event.get("event_type")
        data = event.get("data", {})
        if event_type == "order.created":
            await self._reserve(data)
        elif event_type == "order.cancelled":
            await self._release(data)

    async def _reserve(self, data: dict[str, Any]) -> None:
        for item in data["items"]:
            await self.inventory.reserve(
                item["product_id"], item["quantity"], data["order_id"]
            )

    async def _release(self, data: dict[str, Any]) -> None:
        await self.inventory.release_by_order(data["order_id"])

# kafka/producer.py
import json
import logging
import os
import time
from typing import Any, Optional

from aiokafka import AIOKafkaProducer

logger = logging.getLogger(__name__)


class KafkaProducer:
    """Thin wrapper around AIOKafkaProducer. Uses idempotent + acks=all so that
    retries don't produce duplicates and writes aren't acknowledged until every
    in-sync replica has the record."""

    def __init__(self) -> None:
        self.service_name = os.getenv("SERVICE_NAME", "unknown")
        self._producer: Optional[AIOKafkaProducer] = None

    async def connect(self) -> None:
        self._producer = AIOKafkaProducer(
            bootstrap_servers=os.getenv("KAFKA_BROKERS", "localhost:9092"),
            client_id=self.service_name,
            enable_idempotence=True,   # prevents duplicates on retry
            acks="all",                # wait for all ISR replicas
            compression_type="snappy", # massive throughput win for JSON payloads
            max_batch_size=16384,
            linger_ms=10,              # small batching window for better throughput
        )
        await self._producer.start()
        logger.info("Kafka producer connected")

    async def disconnect(self) -> None:
        if self._producer:
            await self._producer.stop()

    async def send(
        self,
        topic: str,
        message: dict[str, Any],
        key: Optional[str] = None,
    ) -> None:
        assert self._producer is not None
        # Key choice is critical: same key -> same partition -> ordered.
        # Use entity IDs (order_id, user_id) as keys, never timestamps or UUIDs
        # unless you explicitly don't care about ordering.
        headers = [
            ("event-type", message.get("event_type", "").encode()),
            ("event-id", str(message.get("event_id", "")).encode()),
            ("correlation-id", str(message.get("correlation_id", "")).encode()),
            ("source-service", self.service_name.encode()),
            ("timestamp", str(int(time.time() * 1000)).encode()),
        ]
        await self._producer.send_and_wait(
            topic=topic,
            value=json.dumps(message).encode(),
            key=key.encode() if key else None,
            headers=headers,
        )

    async def send_batch(
        self, topic: str, messages: list[dict[str, Any]]
    ) -> None:
        assert self._producer is not None
        # aiokafka batches under the hood based on linger_ms + max_batch_size,
        # so we just fire the sends concurrently.
        batch = self._producer.create_batch()
        for msg in messages:
            batch.append(
                key=msg["key"].encode() if msg.get("key") else None,
                value=json.dumps(msg["value"]).encode(),
                timestamp=None,
            )
        await self._producer.send_batch(batch, topic, partition=None)

    async def send_in_transaction(
        self, messages: list[dict[str, Any]]
    ) -> None:
        """Exactly-once write across multiple topics.
        Requires transactional_id set on the producer constructor."""
        assert self._producer is not None
        async with self._producer.transaction():
            for item in messages:
                await self._producer.send(
                    topic=item["topic"],
                    value=json.dumps(item["value"]).encode(),
                    key=item["key"].encode() if item.get("key") else None,
                )

# kafka/consumer.py
import json
import logging
import os
from typing import Any, Awaitable, Callable, Iterable

from aiokafka import AIOKafkaConsumer, TopicPartition

logger = logging.getLogger(__name__)

MessageHandler = Callable[[dict[str, Any], dict[str, Any]], Awaitable[None]]


class KafkaConsumer:
    def __init__(self, group_id: str) -> None:
        self.group_id = group_id
        self.service_name = os.getenv("SERVICE_NAME", "unknown")
        self._consumer: AIOKafkaConsumer | None = None

    async def connect(self, topics: Iterable[str], from_beginning: bool = False) -> None:
        self._consumer = AIOKafkaConsumer(
            *topics,
            bootstrap_servers=os.getenv("KAFKA_BROKERS", "localhost:9092"),
            group_id=self.group_id,
            client_id=self.service_name,
            # Manual commits so we can ensure at-least-once after successful processing.
            enable_auto_commit=False,
            auto_offset_reset="earliest" if from_beginning else "latest",
            session_timeout_ms=30_000,
            heartbeat_interval_ms=3_000,
            max_poll_interval_ms=300_000,  # give handlers plenty of time before rebalance
        )
        await self._consumer.start()
        logger.info("Kafka consumer %s connected to %s", self.group_id, list(topics))

    async def disconnect(self) -> None:
        if self._consumer:
            await self._consumer.stop()

    async def consume(self, handler: MessageHandler) -> None:
        """Main consume loop. Commits after successful processing (at-least-once).
        Raising from the handler stops the loop -- caller decides retry/DLQ policy."""
        assert self._consumer is not None
        async for msg in self._consumer:
            event = json.loads(msg.value)
            metadata = {
                "topic": msg.topic,
                "partition": msg.partition,
                "offset": msg.offset,
                "key": msg.key.decode() if msg.key else None,
                "timestamp": msg.timestamp,
                "headers": {
                    k: v.decode() if isinstance(v, (bytes, bytearray)) else v
                    for k, v in (msg.headers or [])
                },
            }
            try:
                await handler(event, metadata)
                # Commit AFTER successful processing. Commit the next offset
                # (the position to resume from), which is current + 1.
                tp = TopicPartition(msg.topic, msg.partition)
                await self._consumer.commit({tp: msg.offset + 1})
            except Exception:
                logger.exception("Error processing %s@%d:%d", msg.topic, msg.partition, msg.offset)
                # Caller's responsibility: retry, DLQ, or re-raise.
                raise

    async def consume_batches(
        self,
        handler: Callable[[list[dict[str, Any]], str, int], Awaitable[None]],
        batch_size: int = 100,
        timeout_ms: int = 1000,
    ) -> None:
        """Batch consumption improves throughput when the handler can do
        bulk work (batch DB inserts, bulk external API calls)."""
        assert self._consumer is not None
        while True:
            records = await self._consumer.getmany(
                timeout_ms=timeout_ms, max_records=batch_size
            )
            for tp, messages in records.items():
                if not messages:
                    continue
                events = [json.loads(m.value) for m in messages]
                await handler(events, tp.topic, tp.partition)
                # Commit the offset after the last message in this batch.
                await self._consumer.commit({tp: messages[-1].offset + 1})

# order_service/kafka_events.py
import uuid
from datetime import datetime, timezone
from typing import Any

from kafka.producer import KafkaProducer


class OrderKafkaPublisher:
    def __init__(self) -> None:
        self.producer = KafkaProducer()

    async def initialize(self) -> None:
        await self.producer.connect()

    async def order_created(self, order: dict[str, Any]) -> None:
        # key=order["id"] guarantees all events for this order go to the same
        # partition -> consumers process them in the order they were published.
        await self.producer.send(
            topic="orders",
            message={
                "event_id": f"evt-{uuid.uuid4()}",
                "event_type": "order.created",
                "occurred_at": datetime.now(timezone.utc).isoformat(),
                "data": {
                    "order_id": order["id"],
                    "customer_id": order["customer_id"],
                    "items": order["items"],
                    "total_amount": str(order["total_amount"]),
                },
            },
            key=order["id"],
        )

    async def order_status_changed(
        self, order_id: str, old_status: str, new_status: str
    ) -> None:
        await self.producer.send(
            topic="orders",
            message={
                "event_id": f"evt-{uuid.uuid4()}",
                "event_type": "order.status_changed",
                "data": {
                    "order_id": order_id,
                    "old_status": old_status,
                    "new_status": new_status,
                },
            },
            key=order_id,
        )


# inventory_service/kafka_handler.py
from collections import deque
from typing import Any

from kafka.consumer import KafkaConsumer


class InventoryKafkaHandler:
    def __init__(self, inventory_service) -> None:
        self.consumer = KafkaConsumer(group_id="inventory-service")
        self.inventory = inventory_service
        # Sliding window of recently seen event IDs. Cheap and good enough
        # for single-instance deployments. Swap for Redis in production.
        self._processed: set[str] = set()
        self._processed_order: deque[str] = deque(maxlen=10_000)

    async def start(self) -> None:
        await self.consumer.connect(topics=["orders"], from_beginning=False)
        await self.consumer.consume(self._on_event)

    async def _on_event(self, event: dict[str, Any], metadata: dict[str, Any]) -> None:
        event_id = event["event_id"]
        if event_id in self._processed:
            return  # duplicate

        event_type = event["event_type"]
        if event_type == "order.created":
            await self._handle_created(event["data"])
        elif event_type == "order.cancelled":
            await self._handle_cancelled(event["data"])

        self._processed.add(event_id)
        self._processed_order.append(event_id)
        if len(self._processed) > 10_000:
            # Evict oldest IDs so the set doesn't grow unbounded.
            while len(self._processed) > 10_000:
                old = self._processed_order.popleft()
                self._processed.discard(old)

    async def _handle_created(self, data: dict[str, Any]) -> None:
        for item in data["items"]:
            await self.inventory.reserve(
                item["product_id"], item["quantity"], data["order_id"]
            )

    async def _handle_cancelled(self, data: dict[str, Any]) -> None:
        await self.inventory.release_by_order(data["order_id"])

# Canonical event schema as a Pydantic model.
# Using Pydantic gives you validation at both publish-time (producer) and
# consume-time (consumer) -- a malformed event fails fast instead of silently
# corrupting downstream state.
from datetime import datetime
from decimal import Decimal
from typing import Any, Literal, Optional

from pydantic import BaseModel, Field


class EventSource(BaseModel):
    service: str
    instance: str
    environment: Literal["development", "staging", "production"]


class OrderItem(BaseModel):
    product_id: str
    quantity: int = Field(ge=1)
    unit_price: Decimal


class OrderCreatedData(BaseModel):
    order_id: str
    customer_id: str
    items: list[OrderItem]
    total_amount: Decimal
    currency: str = "USD"
    status: Literal["CREATED", "PAID", "SHIPPED", "CANCELLED"]
    created_at: datetime


class OrderCreatedEvent(BaseModel):
    # Envelope
    event_id: str
    event_type: Literal["order.created"] = "order.created"
    event_version: str = "1.2"
    occurred_at: datetime
    correlation_id: Optional[str] = None
    causation_id: Optional[str] = None
    source: EventSource
    # Payload
    data: OrderCreatedData


# Event type naming convention -- entity.action, lowercase, dot-separated.
EVENT_TYPES: dict[str, str] = {
    # Lifecycle
    "order.created": "When an order is placed",
    "order.paid": "When payment is confirmed",
    "order.shipped": "When order leaves warehouse",
    "order.delivered": "When customer receives order",
    "order.cancelled": "When order is cancelled",
    "order.refunded": "When refund is processed",
    # Generic state change
    "order.status_changed": "Generic status change",
    # Sub-entity events
    "order.item_added": "Item added to existing order",
    "order.item_removed": "Item removed from order",
}