market/trading/models.py

from datetime import datetime
from typing import Mapping, Optional, Sequence
from django.conf import settings
from django.contrib.auth.models import AbstractBaseUser, PermissionsMixin
from django.db import models
from django.db.models import Q, Sum
from django.db.models.signals import post_save, pre_save
from django.dispatch import receiver
from django.urls import reverse
from django.utils.translation import gettext_lazy as _
from frozendict import frozendict
from guardian.shortcuts import assign_perm
from hashid_field import HashidAutoField
from trading.managers import UserManager


class User(PermissionsMixin, AbstractBaseUser):
    """
    Base user model, with auth information already filled in.
    """

    username = models.CharField(max_length=30, unique=True)
    email = models.EmailField(unique=True)
    is_superuser = models.BooleanField(default=False)
    is_active = models.BooleanField(default=True)
    date_joined = models.DateTimeField(_("date joined"), auto_now_add=True)

    max_commodities_allowed = models.PositiveIntegerField(
        default=settings.TRADING_FREE_COMMODITIES,
        null=True,
        verbose_name=_("number of commodities this user is allowed to create"),
    )

    objects = UserManager()

    USERNAME_FIELD = "username"
    EMAIL_FIELD = "email"
    REQUIRED_FIELDS = (EMAIL_FIELD,)

    class Meta:
        verbose_name = _("user")
        verbose_name_plural = _("users")

    def get_full_name(self):
        """
        Returns the first_name plus the last_name, with a space in between.
        """
        return self.username

    def get_short_name(self):
        """
        Returns the short name for the user.
        """
        return self.username

    def email_user(self, subject, message, from_email=None, **kwargs):
        """
        Sends an email to this User.
        """
        send_mail(subject, message, from_email, [self.email], **kwargs)

    def can_create_commodity(self) -> bool:
        if self.max_commodities_allowed is None:
            return True
        return self.max_commodities_allowed > self.commodity_set.count()

    def unused_commodities(self) -> Optional[int]:
        if self.max_commodities_allowed is None:
            return None
        return self.max_commodities_allowed - self.commodity_set.count()

    def unused_invites(self) -> Sequence["Invite"]:
        return self.invite_set.filter(accepted_user__isnull=True)

    def balance_of(self, commodity: "Commodity") -> int:
        return self.balances().get(commodity, 0)

    def balances(self) -> Mapping["Commodity", int]:
        # We are the source, receiving from the destination
        credits_sent = {
            row["dest_sends"]: row["dest_amount__sum"]
            for row in self.tx_sent.filter(status=TxRequest.ACCEPTED)
            .values("dest_sends")
            .annotate(Sum("dest_amount"))
        }
        # We are the destination, accepting from the source
        credits_received = {
            row["source_sends"]: row["source_amount__sum"]
            for row in self.tx_received.filter(status=TxRequest.ACCEPTED)
            .values("source_sends")
            .annotate(Sum("source_amount"))
        }
        # We are the source, sending to the destination
        debits_sent = {
            row["source_sends"]: row["source_amount__sum"]
            for row in self.tx_sent.filter(~Q(status=TxRequest.DECLINED))
            .values("source_sends")
            .annotate(Sum("source_amount"))
        }
        # We are the desination, purchasing from the source
        debits_received = {
            row["dest_sends"]: row["dest_amount__sum"]
            for row in self.tx_received.filter(status=TxRequest.ACCEPTED)
            .values("dest_sends")
            .annotate(Sum("dest_amount"))
        }
        ipos = {
            row["pk"]: row["in_circulation"]
            for row in self.commodity_set.values("pk", "in_circulation")
        }
        keys = (
            set(ipos.keys())
            | set(credits_sent.keys())
            | set(credits_received.keys())
            | set(debits_sent.keys())
            | set(debits_received.keys())
            # | set(tx_requests.keys())
        )
        return frozendict(
            [
                (
                    Commodity.objects.get(pk=pk),
                    credits_sent.get(pk, 0)
                    + credits_received.get(pk, 0)
                    + ipos.get(pk, 0)
                    - debits_sent.get(pk, 0)
                    - debits_received.get(pk, 0)
                    # - tx_requests.get(pk, 0),
                )
                for pk in keys
            ]
        )


@receiver(post_save, sender=User)
def _user_post_save(sender, instance, created, **kwargs):
    # Create invites for the user, if they are permitted invites
    if created and settings.TRADING_FREE_INVITES:
        for _ in range(settings.TRADING_FREE_INVITES):
            Invite.objects.create(inviter=instance).save()


class Invite(models.Model):
    """
    An invitation that is given out to a user.
    """

    id = HashidAutoField(primary_key=True)
    inviter = models.ForeignKey(User, on_delete=models.SET_NULL, null=True)
    accepted_user = models.ForeignKey(
        User,
        on_delete=models.CASCADE,
        null=True,
        default=None,
        related_name="+",
        verbose_name=_("accepted user"),
    )

    class Meta:
        verbose_name = _("invite")
        verbose_name_plural = _("invites")


class Commodity(models.Model):
    """
    A commodity with some kind of market cap.

    A commodity has an associated "owner" account, which determines where the IPO for commodity
    goes.
    """

    id = HashidAutoField(primary_key=True)
    name = models.CharField(
        blank=False,
        unique=True,
        editable=False,
        max_length=100,
        verbose_name=_("commodity name"),
    )
    symbol = models.CharField(
        blank=False, unique=True, editable=False, max_length=6, verbose_name=_("symbol")
    )
    in_circulation = models.PositiveIntegerField()
    created_by = models.ForeignKey(User, null=True, on_delete=models.SET_NULL)

    class Meta:
        verbose_name = _("commodity")
        verbose_name_plural = _("commodities")

    def trading_volume(self, since: Optional[datetime] = None) -> int:
        """
        Gets the amount of this commodity traded since the given cutoff, inclusive.

        If no cutoff is given, then it returns the volume of all transactions.
        """
        source_filt = {"source_sends": self, "status": TxRequest.ACCEPTED}
        dest_filt = {"dest_sends": self, "status": TxRequest.ACCEPTED}
        if since is not None:
            source_filt["update_time__gte"] = since
            dest_filt["update_time__gte"] = since
        source_query = (
            TxRequest.objects.filter(**source_filt)
            .annotate(Sum("source_amount"))
            .values("source_amount__sum")
            .first()
        )
        dest_query = (
            TxRequest.objects.filter(**dest_filt)
            .annotate(Sum("dest_amount"))
            .values("dest_amount__sum")
            .first()
        )
        source_amount = source_query["source_amount__sum"] if source_query else 0
        dest_amount = dest_query["dest_amount__sum"] if dest_query else 0
        return source_amount + dest_amount

    @property
    def market_spread(self) -> int:
        count = 0
        for user in User.objects.all():
            if not user.is_active:
                continue
            if user.balance_of(self) > 0:
                count += 1
        return count

    @property
    def market_spread_ratio(self) -> float:
        user_count = len(User.objects.all())
        if user_count == 0:
            return 0.0
        return self.market_spread / user_count

    @property
    def amount_pending(self) -> int:
        """
        Gets the amount of this commodity that is currently being held in pending transactions.
        """
        pending = (
            TxRequest.objects.filter(source_sends=self, status=TxRequest.OPEN)
            .annotate(Sum("source_amount"))
            .values("source_amount__sum")
            .first()
        )
        return pending["source_amount__sum"] if pending else 0

    @property
    def url(self) -> str:
        return f'<a href="{reverse("trading:commodity_detail", args=[self.id])}">{self.symbol}</a>'


class MaxCommodityError(Exception):
    def __init__(self, user: User):
        self.user = user
        super().__init__("maximum commodity count reached")


class TxRequest(models.Model):
    """
    A transaction request between two users.
    """

    OPEN = "OPEN"
    DECLINED = "DECL"
    # REVISED = "REV"
    ACCEPTED = "ACC"

    STATUS_CHOICES = [
        (OPEN, "Open"),
        (DECLINED, "Declined"),
        # (REVISED, "Revised"),
        (ACCEPTED, "Accepted"),
    ]

    id = HashidAutoField(primary_key=True)
    status = models.CharField(
        max_length=max([len(c) for c in STATUS_CHOICES]),
        choices=STATUS_CHOICES,
        default=OPEN,
    )

    # previous_request = models.OneToOneField(
    #    "TxRequest",
    #    null=True,
    #    default=None,
    #    on_delete=models.SET_NULL,
    #    related_name="revised_request",
    # )

    create_time = models.DateTimeField(auto_now_add=True, editable=False)
    update_time = models.DateTimeField(auto_now=True)

    source = models.ForeignKey(
        User, on_delete=models.CASCADE, related_name="tx_sent", editable=False,
    )
    dest = models.ForeignKey(
        User, on_delete=models.CASCADE, related_name="tx_received", editable=False,
    )

    source_sends = models.ForeignKey(
        Commodity,
        on_delete=models.CASCADE,
        related_name="+",
        null=True,
        editable=False,
    )
    dest_sends = models.ForeignKey(
        Commodity,
        on_delete=models.CASCADE,
        related_name="+",
        null=True,
        editable=False,
    )

    source_amount = models.PositiveIntegerField(editable=False)
    dest_amount = models.PositiveIntegerField(editable=False)

    @staticmethod
    def open(
        source: User,
        dest: User,
        source_sends: Commodity,
        dest_sends: Commodity,
        source_amount: int,
        dest_amount: int,
    ) -> "TxRequest":
        """
        Opens a new transaction request.
        """
        # Check balance
        source_balance = source.balance_of(source_sends)
        if source_balance < source_amount:
            raise BalanceError(
                source_balance, source, dest, source_amount, source_sends
            )
        req = TxRequest.objects.create(
            status=TxRequest.OPEN,
            source=source,
            dest=dest,
            source_sends=source_sends,
            dest_sends=dest_sends,
            source_amount=source_amount,
            dest_amount=dest_amount,
        )
        req.save()
        return req

    def can_accept(self) -> bool:
        """
        Gets whether this request can be accepted.
        """
        dest_balance = self.dest.balance_of(self.dest_sends)
        return self.status == TxRequest.OPEN and dest_balance >= self.dest_amount

    def accept(self):
        """
        Accepts an open transaction request.
        """
        assert self.status == TxRequest.OPEN

        # Ensure destination balance before continuing
        if not self.can_accept():
            dest_balance = self.dest.balance_of(self.dest_sends)
            raise BalanceError(
                dest_balance, self.dest, self.source, self.dest_amount, self.dest_sends
            )

        # Update status
        self.status = TxRequest.ACCEPTED
        self.save()

    def decline(self):
        """
        Declines an open transaction request.
        """
        assert self.status == TxRequest.OPEN
        self.status = TxRequest.DECLINED
        self.save()


class BalanceError(Exception):
    def __init__(
        self, balance: int, source: User, dest: User, amount: int, commodity: Commodity
    ) -> None:
        self.balance = balance
        self.source = source
        self.dest = dest
        self.amount = amount
        self.commodity = commodity
        super().__init__("balance too low to create transaction")


@receiver(pre_save, sender=TxRequest)
def __tx_pre_save(sender, instance, *args, **kwargs):
    """
    Ensures that a transaction cannot be created if the source user doesn't have enough of the given
    commodity in their account.
    """
    if instance.pk is None:
        source_balance = instance.source.balance_of(instance.source_sends)
        if instance.source_amount > source_balance:
            raise BalanceError(
                source_balance,
                instance.source,
                instance.dest,
                instance.source_amount,
                instance.source_sends,
            )
        # Do not check dest balance - could be a side channel for determining another account's
        # funds which is not something we want


@receiver(post_save, sender=TxRequest)
def __tx_post_save(sender, instance, created, **kwargs):
    if not created:
        return
    assign_perm("view_txrequest", instance.source, instance)
    assign_perm("view_txrequest", instance.dest, instance)