from peewee import CharField, Check, IntegerField, Model

from defence360agent.contracts.config import IContactMessageType


class IContactThrottle(Model):
    class Meta:
        db_table = "icontact_throttle"

    message_type = CharField(
        primary_key=True,
        constraints=[
            Check(
                "message_type in {}".format(
                    (
                        str(IContactMessageType.MALWARE_FOUND),
                        str(IContactMessageType.SCAN_NOT_SCHEDULED),
                    )
                )
            )
        ],
    )
    timestamp = IntegerField(default=0)


def migrate(migrator, database, fake=False, **kwargs):
    migrator.create_model(IContactThrottle)


def rollback(migrator, database, fake=False, **kwargs):
    IContactThrottle = migrator.orm["icontact_throttle"]
    migrator.remove_model(IContactThrottle)