import asyncio
import logging

import aio_pika
from aio_pika.patterns import RPC
from pydantic import ValidationError

from rabbit_broker.models import ErrorMessage, ProcessedMessage, UnprocessedMessage
from rabbit_broker.queues.base import AsyncAbstractMessageQueue
from rabbit_broker.schemas import ProcessedBrokerMessage, UnprocessedBrokerMessage

logger = logging.getLogger(__name__)


class AsyncRabbitMessageQueue(AsyncAbstractMessageQueue):
    async def consume(self) -> None:
        logger.info(
            "%s.%s: RPC consumer started",
            self.__class__.__name__,
            self.consume.__name__,
        )
        await asyncio.Future()

    async def post_message(self, data: UnprocessedBrokerMessage, worker: str) -> ProcessedBrokerMessage:
        try:
            UnprocessedMessage(**data)
        except ValidationError as error:
            logger.error(
                "%s.%s: UnprocessedMessage validation failed!: %s",
                self.__class__.__name__,
                self.post_message.__name__,
                str(error),
            )
            return ErrorMessage().generate(message=str(error))
        response = await self.rpc.call(worker, kwargs=dict(data=data))
        try:
            ProcessedMessage(**response)
        except ValidationError as error:
            return ErrorMessage().generate(
                request_id=data["request_id"],
                request_type=data["request_type"],
                dst=data["src"],
                src=data["dst"],
                message=str(error),
            )
        return response

    async def register_tasks(self, routing_key: str, worker: callable):
        """Вызывать перед стартом консьюмера."""
        await self.rpc.register(routing_key, worker, auto_delete=True)

    async def __aenter__(self):
        """
        Метод входа в контекст подключения
        """
        if self.connection is None or self.connection.is_closed:
            logger.info(
                "%s.%s: Created connection",
                self.__class__.__name__,
                self.__aenter__.__name__,
            )
            self.connection = await aio_pika.connect_robust(
                self.broker_url,
            )
            self.channel = await self.connection.channel()
            self.rpc = await RPC.create(self.channel)
        return self

    async def __aexit__(self, *args, **kwargs):
        await self.connection.close()
        await self.channel.close()
