Lotty/src/backend/apps/notifications/models.py

from typing import override

from django.db import models
from django.utils.translation import gettext_lazy as _

from apps.core.models import BaseModel


class ChannelType(models.TextChoices):
    TELEGRAM = "telegram", _("Telegram")
    SMTP = "smtp", _("SMTP Email")


class NotificationEventType(models.TextChoices):
    EXPERIMENT_STARTED = "experiment_started", _("Experiment started")
    EXPERIMENT_PAUSED = "experiment_paused", _("Experiment paused")
    EXPERIMENT_RESUMED = "experiment_resumed", _("Experiment resumed")
    EXPERIMENT_COMPLETED = "experiment_completed", _("Experiment completed")
    GUARDRAIL_TRIGGERED = "guardrail_triggered", _("Guardrail triggered")
    REVIEW_REQUESTED = "review_requested", _("Review requested")
    REVIEW_APPROVED = "review_approved", _("Review approved")
    REVIEW_REJECTED = "review_rejected", _("Review rejected")


class NotificationStatus(models.TextChoices):
    PENDING = "pending", _("Pending")
    SENT = "sent", _("Sent")
    FAILED = "failed", _("Failed")


class NotificationChannel(BaseModel):
    channel_type = models.CharField(
        max_length=20,
        choices=ChannelType.choices,
        verbose_name=_("channel type"),
    )
    name = models.CharField(
        max_length=200,
        verbose_name=_("name"),
    )
    config = models.JSONField(
        default=dict,
        blank=True,
        verbose_name=_("configuration"),
        help_text=_(
            "Provider-specific settings (tokens, chat IDs, SMTP host, etc.)"
        ),
    )
    is_active = models.BooleanField(
        default=True,
        verbose_name=_("is active"),
    )
    created_at = models.DateTimeField(
        auto_now_add=True,
        verbose_name=_("created at"),
    )
    updated_at = models.DateTimeField(
        auto_now=True,
        verbose_name=_("updated at"),
    )

    class Meta:
        verbose_name = _("notification channel")
        verbose_name_plural = _("notification channels")
        ordering = ["-created_at"]

    @override
    def __str__(self) -> str:
        return f"{self.name} ({self.channel_type})"


class NotificationRule(BaseModel):
    event_type = models.CharField(
        max_length=30,
        choices=NotificationEventType.choices,
        verbose_name=_("event type"),
    )
    channel = models.ForeignKey(
        NotificationChannel,
        on_delete=models.CASCADE,
        related_name="rules",
        verbose_name=_("channel"),
    )
    experiment = models.ForeignKey(
        "experiments.Experiment",
        on_delete=models.CASCADE,
        related_name="notification_rules",
        verbose_name=_("experiment"),
        null=True,
        blank=True,
        help_text=_("If null, applies to all experiments."),
    )
    is_active = models.BooleanField(
        default=True,
        verbose_name=_("is active"),
    )
    rate_limit_window_seconds = models.PositiveIntegerField(
        default=60,
        verbose_name=_("rate limit window seconds"),
    )
    rate_limit_max_notifications = models.PositiveIntegerField(
        default=1,
        verbose_name=_("rate limit max notifications"),
    )
    created_at = models.DateTimeField(
        auto_now_add=True,
        verbose_name=_("created at"),
    )
    updated_at = models.DateTimeField(
        auto_now=True,
        verbose_name=_("updated at"),
    )

    class Meta:
        verbose_name = _("notification rule")
        verbose_name_plural = _("notification rules")
        ordering = ["-created_at"]
        indexes = [
            models.Index(
                fields=["event_type", "is_active"],
                name="idx_rule_event_active",
            ),
        ]

    @override
    def __str__(self) -> str:
        scope = self.experiment.name if self.experiment else "all"
        return f"{self.event_type} -> {self.channel.name} ({scope})"


class NotificationLog(BaseModel):
    rule = models.ForeignKey(
        NotificationRule,
        on_delete=models.SET_NULL,
        null=True,
        related_name="logs",
        verbose_name=_("rule"),
    )
    channel = models.ForeignKey(
        NotificationChannel,
        on_delete=models.SET_NULL,
        null=True,
        related_name="logs",
        verbose_name=_("channel"),
    )
    event_type = models.CharField(
        max_length=30,
        choices=NotificationEventType.choices,
        verbose_name=_("event type"),
    )
    event_key = models.CharField(
        max_length=255,
        verbose_name=_("dedup key"),
        db_index=True,
    )
    payload = models.JSONField(
        default=dict,
        verbose_name=_("payload"),
    )
    status = models.CharField(
        max_length=10,
        choices=NotificationStatus.choices,
        default=NotificationStatus.PENDING,
        verbose_name=_("status"),
    )
    error = models.TextField(
        blank=True,
        default="",
        verbose_name=_("error details"),
    )
    created_at = models.DateTimeField(
        auto_now_add=True,
        verbose_name=_("created at"),
    )
    sent_at = models.DateTimeField(
        null=True,
        blank=True,
        verbose_name=_("sent at"),
    )

    class Meta:
        verbose_name = _("notification log")
        verbose_name_plural = _("notification logs")
        ordering = ["-created_at"]
        indexes = [
            models.Index(
                fields=["status", "-created_at"],
                name="idx_notif_log_status",
            ),
        ]

    @override
    def __str__(self) -> str:
        return f"Notification({self.event_type}, {self.status})"