from __future__ import annotations

import asyncio
import base64
import time
from typing import Any

from aio_pika import Message, connect_robust
from aio_pika.abc import AbstractIncomingMessage

from .config import RabbitMQConfig
from .mcp_errors import MCPToolError


class AMQPManager:
    _connection = None
    _channel: Any = None
    _lock = asyncio.Lock()
    _messages: dict[str, AbstractIncomingMessage] = {}
    _timestamps: dict[str, float] = {}
    _ttl_seconds = 300.0

    @classmethod
    async def _ensure_channel(cls, config: RabbitMQConfig) -> Any:
        async with cls._lock:
            if cls._channel and not cls._channel.is_closed:
                return cls._channel
            cls._connection = await connect_robust(config.amqp_url)
            cls._channel = await cls._connection.channel()
            return cls._channel

    @classmethod
    def _receipt_handle(cls, message: AbstractIncomingMessage) -> str:
        channel_number = getattr(message.channel, "channel_number", 0)
        return f"{channel_number}:{message.delivery_tag}"

    @classmethod
    def _cleanup(cls) -> None:
        now = time.monotonic()
        expired = [
            handle
            for handle, ts in cls._timestamps.items()
            if now - ts > cls._ttl_seconds
        ]
        for handle in expired:
            cls._messages.pop(handle, None)
            cls._timestamps.pop(handle, None)

    @classmethod
    def store_message(cls, message: AbstractIncomingMessage) -> str:
        cls._cleanup()
        handle = cls._receipt_handle(message)
        cls._messages[handle] = message
        cls._timestamps[handle] = time.monotonic()
        return handle

    @classmethod
    def pop_message(cls, handle: str) -> AbstractIncomingMessage:
        message = cls._messages.pop(handle, None)
        cls._timestamps.pop(handle, None)
        if message is None:
            raise MCPToolError(
                code="UNKNOWN_RECEIPT",
                message="Unknown receipt_handle",
                details={"receipt_handle": handle},
            )
        return message


def _encode_payload(payload: str, payload_encoding: str) -> bytes:
    if payload_encoding == "base64":
        return base64.b64decode(payload)
    return payload.encode("utf-8")


def _decode_payload(body: bytes) -> str:
    try:
        return body.decode("utf-8")
    except UnicodeDecodeError:
        return base64.b64encode(body).decode("utf-8")


async def amqp_publish(
    config: RabbitMQConfig,
    exchange: str,
    routing_key: str,
    payload: str,
    payload_encoding: str = "string",
    properties: dict[str, Any] | None = None,
) -> dict[str, Any]:
    channel = await AMQPManager._ensure_channel(config)
    message = Message(body=_encode_payload(payload, payload_encoding), headers=properties)
    if exchange in {"", "amq.default"}:
        await channel.default_exchange.publish(message, routing_key=routing_key)
    else:
        exchange_obj = await channel.get_exchange(exchange, ensure=False)
        await exchange_obj.publish(message, routing_key=routing_key)
    return {"published": True}


async def amqp_consume(
    config: RabbitMQConfig,
    queue: str,
    max_messages: int = 1,
    auto_ack: bool = True,
) -> dict[str, Any]:
    channel = await AMQPManager._ensure_channel(config)
    queue_obj = await channel.declare_queue(queue, passive=True)
    items: list[dict[str, Any]] = []
    for _ in range(max_messages):
        message = await queue_obj.get(no_ack=auto_ack, fail=False)
        if message is None:
            break
        receipt_handle = None
        if not auto_ack:
            receipt_handle = AMQPManager.store_message(message)
        items.append(
            {
                "receipt_handle": receipt_handle,
                "delivery_tag": message.delivery_tag,
                "exchange": message.exchange,
                "routing_key": message.routing_key,
                "redelivered": message.redelivered,
                "payload": _decode_payload(message.body),
            }
        )
    return {"messages": items}


async def amqp_ack(config: RabbitMQConfig, receipt_handle: str) -> dict[str, Any]:
    message = AMQPManager.pop_message(receipt_handle)
    await message.ack()
    return {"acknowledged": True}


async def amqp_nack(
    config: RabbitMQConfig, receipt_handle: str, requeue: bool = True
) -> dict[str, Any]:
    message = AMQPManager.pop_message(receipt_handle)
    await message.nack(requeue=requeue)
    return {"nacked": True, "requeue": requeue}


async def amqp_reject(
    config: RabbitMQConfig, receipt_handle: str, requeue: bool = False
) -> dict[str, Any]:
    message = AMQPManager.pop_message(receipt_handle)
    await message.reject(requeue=requeue)
    return {"rejected": True, "requeue": requeue}
