import asyncio
import logging
import uuid
from asyncio import CancelledError
from collections import deque
from contextlib import (
    AbstractAsyncContextManager,
    asynccontextmanager,
    nullcontext,
    suppress,
)
from datetime import UTC, datetime
from typing import Any, AsyncIterator, Deque, Mapping, Protocol, cast

import pymongo.errors
from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
from bson import ObjectId, Timestamp
from motor.motor_asyncio import (
    AsyncIOMotorClient,
    AsyncIOMotorClientSession,
    AsyncIOMotorDatabase,
)
from pydantic import AwareDatetime, BaseModel, ConfigDict, Field

__all__ = ["Event", "EventListener", "EventOutbox", "EventHandler"]


class Event(BaseModel):
    topic: str
    content_schema: str
    idempotency_key: str = Field(default_factory=lambda: uuid.uuid4().hex)
    occurred_at: AwareDatetime = Field(default_factory=lambda: datetime.now(tz=UTC))

    model_config = ConfigDict(extra="allow")


class EventListener:
    def __init__(self, events: Deque[Event]) -> None:
        self.__events = events

    def event_occurred(self, event: Event) -> None:
        self.__events.append(event)


class EventHandler(Protocol):
    async def __call__(
        self,
        event: Event,
        mongo_session: AsyncIOMotorClientSession,
        /,
    ) -> None:
        pass  # pragma: no cover


class EventOutbox:
    def __init__(
        self,
        mongo_client: AsyncIOMotorClient,
        kafka_producer: AIOKafkaProducer,
        kafka_consumer: AIOKafkaConsumer,
        *,
        mongo_db: AsyncIOMotorDatabase | None = None,
        mongo_collection_outbox: str = "transactional-outbox",
        mongo_collection_inbox: str = "transactional-inbox",
    ) -> None:
        self.mongo_client = mongo_client
        self.mongo_db = (
            mongo_db
            if mongo_db is not None
            else self.mongo_client.get_default_database()
        )
        self.mongo_collection_outbox = mongo_collection_outbox
        self.mongo_collection_inbox = mongo_collection_inbox
        self.kafka_consumer = kafka_consumer
        self.kafka_producer = kafka_producer

    async def create_indexes(self) -> None:
        await self._outbox.create_index(
            [
                ("payload.topic", 1),
                ("payload.content_schema", 1),
                ("payload.idempotency_key", 1),
            ],
            name="idempotency",
            unique=True,
        )

    def event_listener(
        self, mongo_session: AsyncIOMotorClientSession
    ) -> AbstractAsyncContextManager[EventListener]:
        async def event_listener() -> AsyncIterator[EventListener]:
            events: Deque[Event] = deque()
            listener = EventListener(events)
            async with _ensure_session_in_transaction(mongo_session):
                yield listener
                documents = []
                while events:
                    event = events.popleft()
                    document = {
                        "payload": event.model_dump(mode="json"),
                        "published": False,
                    }
                    documents.append(document)
                if documents:
                    await self._outbox.insert_many(
                        documents,
                        session=mongo_session,
                    )

        return asynccontextmanager(event_listener)()

    def run_event_handler(
        self, handler: EventHandler
    ) -> AbstractAsyncContextManager[None]:
        async def run_event_handler() -> AsyncIterator[None]:
            async with await self.mongo_client.start_session() as mongo_session:
                producer_loop = asyncio.create_task(
                    self._produce_messages(mongo_session),
                )
                consumer_loop = asyncio.create_task(
                    self._consume_messages(handler),
                )
                try:
                    yield
                finally:
                    consumer_loop.cancel()
                    producer_loop.cancel()
                    with suppress(CancelledError):
                        await asyncio.gather(consumer_loop, producer_loop)

        return asynccontextmanager(run_event_handler)()

    @property
    def _outbox(self):
        return self.mongo_db[self.mongo_collection_outbox]

    @property
    def _inbox(self):
        return self.mongo_db[self.mongo_collection_inbox]

    async def _produce_messages(self, mongo_session: AsyncIOMotorClientSession) -> None:
        latest_handled_object_id: ObjectId | None = None
        async for document in self._outbox.find(
            {"published": False},
            session=mongo_session,
        ):
            await self._publish_event(mongo_session, document)
            latest_handled_object_id = document["_id"]

        if latest_handled_object_id:
            id_generated_at = latest_handled_object_id.generation_time
            start_at_operation_time = Timestamp(id_generated_at, 0)
        else:
            cluster_time_document = cast(Mapping[str, Any], mongo_session.cluster_time)
            start_at_operation_time = cluster_time_document["clusterTime"]

        pipeline = [{"$match": {"operationType": {"$in": ["insert", "invalidate"]}}}]
        change_stream = self._outbox.watch(
            pipeline,
            start_at_operation_time=start_at_operation_time,
            session=mongo_session,
        )
        while True:
            try:
                async with change_stream:
                    async for change_event in change_stream:
                        if change_event["operationType"] == "invalidate":
                            raise _ChangeStreamInvalidated(change_event["_id"])
                        document = change_event["fullDocument"]
                        await self._publish_event(mongo_session, document)
            except _ChangeStreamInvalidated as ex:
                change_stream = self._outbox.watch(
                    pipeline,
                    start_after=ex.resume_token,
                    session=mongo_session,
                )

    async def _publish_event(
        self,
        mongo_session: AsyncIOMotorClientSession,
        document: Mapping[str, Any],
    ) -> None:
        event = Event.model_validate(document["payload"])
        async with mongo_session.start_transaction():
            published_at = datetime.now(tz=UTC)
            update_result = await self._outbox.update_one(
                {"_id": document["_id"], "published": False},
                {"$set": {"published": True, "published_at": published_at}},
                session=mongo_session,
            )
            if update_result.modified_count:
                await self.kafka_producer.send_and_wait(
                    event.topic,
                    event.model_dump_json().encode(),
                    partition=0,
                    timestamp_ms=int(published_at.timestamp()),
                )
            else:
                logging.getLogger(__name__).info(
                    "Already published event %r from %r collection skipped",
                    document["_id"],
                    self.mongo_collection_outbox,
                )

    async def _consume_messages(self, handler: EventHandler) -> None:
        async with await self.mongo_client.start_session() as mongo_session:
            while True:
                kafka_consumer_record = await self.kafka_consumer.getone()
                event = Event.model_validate_json(kafka_consumer_record.value)
                try:
                    await self._inbox.insert_one(
                        {
                            "_id": event.model_dump(
                                mode="json",
                                include={"topic", "content_schema", "idempotency_key"},
                            ),
                            "handled": False,
                        },
                        session=mongo_session,
                    )
                except pymongo.errors.DuplicateKeyError:
                    pass

                while True:
                    try:
                        await self._handle_event(event, mongo_session, handler)
                    except Exception as ex:
                        logging.getLogger(__name__).critical(
                            "Failed to handle event %r: %r",
                            event.model_dump(
                                mode="json",
                                include={"topic", "content_schema", "idempotency_key"},
                            ),
                            ex,
                        )
                        continue
                    await self.kafka_consumer.commit()
                    break

    async def _handle_event(
        self,
        event: Event,
        mongo_session: AsyncIOMotorClientSession,
        handler: EventHandler,
    ) -> None:
        async with mongo_session.start_transaction():
            result = await self._inbox.update_one(
                {
                    "_id": event.model_dump(
                        mode="json",
                        include={
                            "topic",
                            "content_schema",
                            "idempotency_key",
                        },
                    ),
                    "handled": False,
                },
                {
                    "$set": {
                        "handled": True,
                        "handled_at": datetime.now(tz=UTC),
                    }
                },
                session=mongo_session,
            )
            if not result.modified_count:
                return

            await handler(event, mongo_session)


def _ensure_session_in_transaction(
    mongo_session: AsyncIOMotorClientSession,
) -> AbstractAsyncContextManager[Any]:
    if bool(mongo_session.in_transaction):
        return nullcontext()
    return mongo_session.start_transaction()


class _ChangeStreamInvalidated(Exception):
    def __init__(self, resume_token: Mapping[str, Any]) -> None:
        self.resume_token = resume_token
