#!/usr/bin/env python3
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#

from __future__ import annotations

from abc import ABC, abstractmethod
from collections import deque, namedtuple
from dataclasses import dataclass
import optparse
import time
import uuid

from typing import Optional, Union

from proton import (Condition, Delivery, Described, Disposition, DispositionType,
                    Endpoint, Link, Sender, Message, Terminus, DeclaredDisposition,
                    RejectedDisposition, TransactionalDisposition)
from proton.handlers import MessagingHandler
from proton.reactor import Container


class Queue:
    def __init__(self, broker: Broker, name: str, dynamic: bool = False):
        self.broker: Broker = broker
        self.name: str = name
        self.dynamic: bool = dynamic
        self.queue: deque[Message] = deque()
        self.consumers: list[Sender] = []

    def subscribe(self, consumer: Sender):
        self.consumers.append(consumer)

    def unsubscribe(self, consumer: Sender):
        if consumer in self.consumers:
            self.consumers.remove(consumer)

    def removable(self):
        return len(self.consumers) == 0 and (self.dynamic or len(self.queue) == 0)

    def publish(self, message: Message):
        self.queue.append(message)
        self.dispatch()

    def dispatch(self, consumer: Optional[Sender] = None):
        if consumer:
            c = [consumer]
        else:
            c = self.consumers
        while self._deliver_to(c):
            pass

    def _deliver_to(self, consumers: list[Sender]):
        try:
            result = False
            for c in consumers:
                if c.credit:
                    message = self.queue.popleft()
                    self.broker.deliver(c, message)
                    result = True
            return result
        except IndexError:  # no more messages
            return False


UnsettledDelivery = namedtuple('UnsettledDelivery', ['address', 'message'])


class TransactionAction(ABC):
    @abstractmethod
    def commit(self, broker: Broker) -> None: ...

    @abstractmethod
    def rollback(self, broker: Broker) -> None: ...


@dataclass
class QueueMessage(TransactionAction):
    delivery: Delivery
    message: Message
    address: str

    def commit(self, broker: Broker) -> None:
        broker.publish(self.message, self.address)
        self.delivery.settle()

    def rollback(self, broker: Broker) -> None:
        pass


@dataclass
class RejectDelivery(TransactionAction):
    delivery: Delivery

    def commit(self, broker: Broker) -> None:
        self.delivery.settle()

    def rollback(self, broker: Broker) -> None:
        pass


@dataclass
class DeliveryUpdate(TransactionAction):
    delivery: Delivery
    commit_outcome: DispositionType

    def commit(self, broker: Broker) -> None:
        broker.delivery_outcome(self.delivery, self.commit_outcome)

    def rollback(self, broker: Broker) -> None:
        pass


class Transaction:
    """
    Simple container for broker-managed transactions.
    """
    def __init__(
        self,
        tid: bytes,
        start_time: float,
    ):
        self.tid = tid
        self.start_time = start_time
        self.actions: list[TransactionAction] = []

    def add_action(self, action: TransactionAction) -> None:
        self.actions.append(action)

    def commit(self, broker: Broker) -> Optional[Condition]:
        for action in self.actions:
            action.commit(broker)
        return None

    def rollback(self, broker: Broker) -> None:
        for action in self.actions:
            action.rollback(broker)


class Broker(MessagingHandler):
    def __init__(self, url: str, verbose: bool, redelivery_limit: int, txn_timeout: float) -> None:
        super().__init__(auto_accept=False)
        self.verbose = verbose
        self.url = url
        self.redelivery_limit = redelivery_limit
        self.txn_timeout = txn_timeout
        self.queues: dict[str, Queue] = {}
        self.unsettled_deliveries: dict[Delivery, UnsettledDelivery] = {}
        self.txns: dict[bytes, Transaction] = {}
        self.acceptor = None

    def _verbose_print(self, message):
        if self.verbose:
            print(message)

    def _queue(self, address, dynamic=False):
        if address not in self.queues:
            self.queues[address] = Queue(self, address, dynamic)
            self._verbose_print(f"{address=}: Created")
        return self.queues[address]

    def on_start(self, event):
        self.acceptor = event.container.listen(self.url)

    def on_connection_opening(self, event):
        event.connection.offered_capabilities = 'ANONYMOUS-RELAY'

    def on_link_opening(self, event):
        link = event.link
        if link.is_sender:
            dynamic = link.remote_source.dynamic
            if dynamic or link.remote_source.address:
                address = str(uuid.uuid4()) if dynamic else link.remote_source.address
                link.source.address = address
                self._queue(address, dynamic).subscribe(link)
                self._verbose_print(f"{link=}: Subscribed: {address=}")
        elif link.remote_target.type == Terminus.COORDINATOR:
            # Set up transaction coordinator
            # Should check for compatible capabilities
            # requested = link.remote_target.capabilities.get_object()
            link.target.type = Terminus.COORDINATOR
            link.target.copy(link.remote_target)
            link._txns: set[Transaction] = set()
        elif link.remote_target.address:
            link.target.address = link.remote_target.address

    def _unsubscribe(self, link):
        address = link.source.address
        if address in self.queues:
            q = self.queues[address]
            q.unsubscribe(link)
            self._verbose_print(f"{link=}: Unsubscribed: {address=}")
            if q.removable():
                del q
                self._verbose_print(f"{address=}: Removed")

    def _modify_delivery(self, delivery: Delivery, undeliverable: bool, failed: bool, message: Message, address: str):
        # If not deliverable don't requeue
        if undeliverable:
            self._verbose_print(f"{delivery.tag=}: Modified: Undeliverable: {message.id=}")
            # Don't requeue the message
            return
        # Check if we need to update the delivery count
        if failed:
            if message.delivery_count >= self.redelivery_limit:
                self._verbose_print(f"{delivery.tag=}: Modified: Redelivery limit exceeded: {message.id=}")
                # Don't requeue the message
                return
            # Update the delivery count
            message.delivery_count += 1
        # Requeue the message from the delivery
        self._verbose_print(f"{delivery.tag=}: Modified: {message.id=} Requeued: {address=}")
        self._queue(address).publish(message)

    def delivery_outcome(self, delivery: Delivery, outcome: Union[int, DispositionType]):
        unsettled_delivery = self.unsettled_deliveries[delivery]
        message = unsettled_delivery.message
        address = unsettled_delivery.address
        if outcome == Disposition.ACCEPTED:
            self._verbose_print(f"{delivery.tag=}: Accepted: {message.id=}")
            # Delivery was accepted - nothing further to do
        elif outcome == Disposition.REJECTED:
            self._verbose_print(f"{delivery.tag=}: Rejected: {message.id=}")
            # Delivery was rejected - nothing further to do
        elif outcome == Disposition.RELEASED:
            self._verbose_print(f"{delivery.tag=}: Released: {message.id=} Requeued: {address=}")
            # Requeue the message from the delivery
            self._queue(address).publish(message)
        elif outcome == Disposition.MODIFIED:
            disposition = delivery.remote
            # This is a little hack which assumes that if the outcome is modified but the disposition is not,
            # then the disposition was a transacted modified disposition and should be treated as failed delivery.
            if disposition.type == Disposition.MODIFIED:
                undeliverable = disposition.undeliverable
                failed = disposition.failed
            else:
                undeliverable = False
                failed = True
            self._modify_delivery(delivery, undeliverable=undeliverable, failed=failed, message=message, address=address)
        delivery.settle()
        del unsettled_delivery

    def _declare_txn(self) -> Transaction:
        tid = bytes(uuid.uuid4().bytes)
        txn = Transaction(tid, time.monotonic())
        self.txns[tid] = txn
        return txn

    def _discharge_txn(self, txn: Transaction, failed) -> Optional[Condition]:
        error = None
        tid = txn.tid
        if not failed:
            # Commit
            txn_time = time.monotonic() - txn.start_time
            if self.txn_timeout > 0 and txn_time > self.txn_timeout:
                error = Condition('amqp:transaction:timeout', f"timeout: {self.txn_timeout:.1f}s time: {txn_time:.1f}s")
            else:
                error = txn.commit(self)
            if error is not None:
                self._verbose_print(f"{tid=}: Commit failed: {error} -> rollback")
                txn.rollback(self)
        else:
            # Rollback
            self._verbose_print(f"{tid=}: Rollback")
            txn.rollback(self)
        del self.txns[tid]
        return error

    def _coordinator_message(self, msg, delivery):
        body = msg.body
        if isinstance(body, Described):
            link = delivery.link
            d = body.descriptor
            if d == "amqp:declare:list":
                # Allocate transaction id
                txn = self._declare_txn()
                self._verbose_print(f"{txn.tid=}: Declare")
                delivery.local = DeclaredDisposition(txn.tid)
                link._txns.add(txn)
            elif d == "amqp:discharge:list":
                # Commit/abort the transaction (commit may fail and be rejected).
                value = body.value
                tid = bytes(value[0])
                failed = bool(value[1])
                txn = self.txns.get(tid)
                if txn is not None:
                    error = self._discharge_txn(txn, failed)
                    if error is None:
                        delivery.update(Disposition.ACCEPTED)
                    else:
                        delivery.update(RejectedDisposition(error))
                    link._txns.remove(txn)
                else:
                    self._verbose_print(f"{tid=}: Discharge unknown txn-id: {failed=}")
                    delivery.update(RejectedDisposition(Condition('amqp:transaction:unknown-id')))
        delivery.settle()

    def on_link_closing(self, event):
        link = event.link
        if link.is_sender:
            self._unsubscribe(link)
        elif link.target.type == Terminus.COORDINATOR:
            # Abort any remaining active transactions
            for txn in link._txns:
                self._discharge_txn(txn, failed=True)
            link._txns.clear()

    def _remove_stale_consumers(self, connection):
        for link in connection.links(Endpoint.REMOTE_ACTIVE):
            if link.is_sender:
                self._unsubscribe(link)

    def _abort_active_transactions(self, connection):
        for link in connection.links(Endpoint.LOCAL_ACTIVE):
            if link.target.type == Terminus.COORDINATOR:
                # Abort any remaining active transactions
                for txn in link._txns:
                    self._discharge_txn(txn, failed=True)
                link._txns.clear()

    def on_connection_closing(self, event):
        connection = event.connection
        self._remove_stale_consumers(connection)
        self._abort_active_transactions(connection)

    def on_disconnected(self, event):
        connection = event.connection
        self._remove_stale_consumers(connection)
        self._abort_active_transactions(connection)

    def on_sendable(self, event):
        link: Link = event.link
        self._queue(link.source.address).dispatch(link)

    def on_message(self, event):
        link: Link = event.link
        delivery: Delivery = event.delivery

        message = event.message
        if link.target.type == Terminus.COORDINATOR:
            # Deal with special transaction messages
            self._coordinator_message(message, delivery)
            return

        address = link.target.address
        if address is None:
            address = message.address

        ldisposition = None
        if address is None:
            self._verbose_print(f"{message.id=}: Message without address")
            ldisposition = RejectedDisposition(Condition('amqp:link:invalid-address'))
        else:
            ldisposition = Disposition.ACCEPTED

        # Is this a transactioned message?
        if delivery.remote_state == Disposition.TRANSACTIONAL_STATE:
            rdisposition = delivery.remote
            tid = rdisposition.id
            if tid in self.txns:
                if address:
                    self._verbose_print(f"{tid=}: Message: {message.id=}")
                    self.txns[tid].add_action(QueueMessage(delivery, message, address))
                else:
                    self.txns[tid].add_action(RejectDelivery(delivery))
                delivery.update(TransactionalDisposition(tid, ldisposition))
                return
            else:
                self._verbose_print(f"{tid=}: Message: unknown txn-id")
                delivery.update(RejectedDisposition(Condition('amqp:transaction:unknown-id')))
                delivery.settle()
                return

        if address:
            self.publish(message, address)
        delivery.update(ldisposition)
        delivery.settle()

    def publish(self, message, address):
        queue = self._queue(address)
        queue.publish(message)
        self._verbose_print(f"{message.id=} Queued: {queue.name=}")

    def deliver(self, consumer, message):
        delivery = consumer.send(message)
        address = consumer.source.address
        self.unsettled_deliveries[delivery] = UnsettledDelivery(address, message)
        self._verbose_print(f"{delivery.tag=}: Sent: {message.id=} to {address=}")

    def on_delivery_updated(self, event):
        """Handle all delivery updates for the link."""
        delivery = event.delivery
        disposition_type = delivery.remote_state
        # Is this a transactioned delivery update?
        if disposition_type == Disposition.TRANSACTIONAL_STATE:
            disposition = delivery.remote
            tid = disposition.id
            outcome = disposition.outcome_type
            if tid in self.txns:
                self._verbose_print(f"{tid=}: Delivery update: outcome={outcome}")
                self.txns[tid].add_action(DeliveryUpdate(delivery, outcome))
                return
            else:
                self._verbose_print(f"{tid=}: Delivery update: unknown txn-id")
                delivery.update(RejectedDisposition(Condition('amqp:transaction:unknown-id')))
        else:
            self.delivery_outcome(delivery, disposition_type)
        # The delivery is settled in every case except a valid transaction
        # where the outcome is not yet known until the transaction is discharged.
        delivery.settle()


def main():
    parser = optparse.OptionParser(usage="usage: %prog [options]")
    parser.add_option("-v", "--verbose", action="store_true", default=False,
                      help="enable verbose output (default %default)")
    parser.add_option("-a", "--address", default="localhost:5672",
                      help="address router listens on (default %default)")
    parser.add_option("-l", "--redelivery-limit", default=5,
                      help="maximum redelivery attempts (default %default)")
    parser.add_option("-t", "--txn-timeout", default=0,
                      help="transaction timeout in seconds (default %default)")
    opts, args = parser.parse_args()

    try:
        Container(Broker(opts.address, opts.verbose, int(opts.redelivery_limit), float(opts.txn_timeout))).run()
    except KeyboardInterrupt:
        pass


if __name__ == '__main__':
    main()