lglog/lg/models.py

from django.contrib.auth.models import AbstractUser
from django.db import models
from django.utils.translation import gettext_lazy as _
from polymorphic.models import PolymorphicModel


class DiscordUser(AbstractUser):
    discord_id = models.CharField(
        max_length=32,
        primary_key=True,
        verbose_name=_("discord id"),
    )

    discriminator = models.CharField(
        max_length=4,
        verbose_name=_("discriminator"),
    )

    avatar_id = models.CharField(
        max_length=32,
        verbose_name=_("avatar id"),
    )

    @property
    def avatar_url(self):
        return f"https://cdn.discordapp.com/avatars/{self.discord_id}/{self.avatar_id}.png"


class Player(PolymorphicModel):
    user = models.OneToOneField(
        DiscordUser,
        on_delete=models.CASCADE,
        verbose_name=_("user"),
    )

    dead = models.BooleanField(
        default=False,
        verbose_name=_("dead"),
    )

    infected = models.BooleanField(
        default=False,
        verbose_name=_("infected"),
    )

    hacked = models.BooleanField(
        default=False,
        verbose_name=_("hacked"),
    )

    captain = models.BooleanField(
        default=False,
        verbose_name=_("captain"),
    )

    in_love_with = models.ForeignKey(
        "lg.Player",
        on_delete=models.CASCADE,
        related_name="+",
        null=True,
        default=None,
    )

    @property
    def alive(self):
        return not self.dead

    def action_class(self):
        raise NotImplementedError

    class Meta:
        verbose_name = _("player")
        verbose_name_plural = _("players")


class Werewolf(Player):
    def action_class(self):
        return WerewolfAction

    class Meta:
        verbose_name = _("werewolf")
        verbose_name_plural = _("werewolves")


class Cupidon(Player):
    def action_class(self):
        return CupidonAction

    class Meta:
        verbose_name = _("cupidon")
        verbose_name_plural = _("cupidons")


class Witch(Player):
    def action_class(self):
        return WitchAction

    class Meta:
        verbose_name = _("witch")
        verbose_name_plural = _("witches")


class Raven(Player):
    def action_class(self):
        return RavenAction

    class Meta:
        verbose_name = _("raven")
        verbose_name_plural = _("ravens")


class Dove(Player):
    def action_class(self):
        return DoveAction

    class Meta:
        verbose_name = _("dove")
        verbose_name_plural = _("doves")


class Hacker(Player):
    def action_class(self):
        return HackerAction

    class Meta:
        verbose_name = _("hacker")
        verbose_name_plural = _("hackers")


class Action(PolymorphicModel):
    player = models.ForeignKey(
        Player,
        on_delete=models.CASCADE,
        verbose_name=_("player"),
        related_name="actions",
    )

    night = models.PositiveSmallIntegerField(
        verbose_name=_("night"),
    )

    def act_night(self, night: int) -> None:
        pass

    class Meta:
        verbose_name = _("action")
        verbose_name_plural = _("actions")


class WerewolfAction(Action):
    target = models.ForeignKey(
        Player,
        on_delete=models.CASCADE,
        related_name="+",
        verbose_name=_("target"),
    )

    class Meta:
        verbose_name = _("werewolf action")
        verbose_name_plural = _("werewolf actions")


class CupidonAction(Action):
    target1 = models.ForeignKey(
        Player,
        on_delete=models.CASCADE,
        related_name="+",
        verbose_name=_("target"),
    )

    target2 = models.ForeignKey(
        Player,
        on_delete=models.CASCADE,
        related_name="+",
        verbose_name=_("target"),
    )

    def act_night(self, night: int) -> None:
        if night > 1:
            return
        self.target1.in_love_with = self.target2
        self.target2.in_love_with = self.target1
        self.target1.save()
        self.target2.save()


class WitchAction(Action):
    saved = models.ForeignKey(
        Player,
        on_delete=models.CASCADE,
        related_name="+",
        verbose_name=_("saved"),
    )

    killed = models.ForeignKey(
        Player,
        on_delete=models.CASCADE,
        related_name="+",
        verbose_name=_("killed"),
    )

    class Meta:
        verbose_name = _("witch action")
        verbose_name_plural = _("witch actions")


class RavenAction(Action):
    target = models.ForeignKey(
        Player,
        on_delete=models.CASCADE,
        related_name="+",
        verbose_name=_("target"),
    )

    message = models.TextField(
        verbose_name=_("message"),
    )

    class Meta:
        verbose_name = _("raven action")
        verbose_name_plural = _("raven actions")


class DoveAction(Action):
    target = models.ForeignKey(
        Player,
        on_delete=models.CASCADE,
        related_name="+",
        verbose_name=_("target"),
    )

    message = models.TextField(
        verbose_name=_("message"),
    )

    class Meta:
        verbose_name = _("dove action")
        verbose_name_plural = _("dove actions")


class HackerAction(Action):
    target1 = models.ForeignKey(
        Player,
        on_delete=models.CASCADE,
        related_name="+",
        verbose_name=_("target"),
    )

    target2 = models.ForeignKey(
        Player,
        on_delete=models.CASCADE,
        related_name="+",
        verbose_name=_("target"),
    )

    def act_night(self, night: int):
        self.target1.hacked = True
        self.target2.hacked = True
        self.target1.save()
        self.target2.save()

    class Meta:
        verbose_name = _("hacker action")
        verbose_name_plural = _("hacker actions")


class Vote(PolymorphicModel):
    day = models.PositiveSmallIntegerField(
        verbose_name=_("day"),
    )

    voter = models.ForeignKey(
        Player,
        on_delete=models.CASCADE,
        related_name="+",
        verbose_name=_("voter"),
    )

    target = models.ForeignKey(
        Player,
        on_delete=models.CASCADE,
        related_name="+",
        verbose_name=_("target"),
    )

    @classmethod
    def get_votes(cls, day: int):
        votes = cls.objects.filter(day=day).all()
        d = {}
        for vote in votes:
            d.setdefault(vote.target, 0)
            d[vote.target] += 1
        return d

    class Meta:
        verbose_name = _("vote")
        verbose_name_plural = _("votes")


class CaptainVote(Vote):
    class Meta:
        verbose_name = _("captain vote")
        verbose_name_plural = _("captain votes")


class EjectVote(Vote):
    @classmethod
    def get_votes(cls, day: int):
        votes = cls.objects.filter(day=day).all()
        d = {}
        for vote in votes:
            vote: Vote
            d.setdefault(vote.target, 0)
            d[vote.target] += 1
            if vote.voter.captain:
                d[vote.target] += 1

        for dove_action in DoveAction.objects.filter(day=day).all():
            dove_action: DoveAction
            d.setdefault(dove_action.target, 0)
            d[dove_action.target] -= 2

        for raven_action in RavenAction.objects.filter(day=day).all():
            raven_action: RavenAction
            d.setdefault(raven_action.target, 0)
            d[raven_action.target] += 2

        return d

    class Meta:
        verbose_name = _("eject vote")
        verbose_name_plural = _("eject votes")