main-site/apps/blog/models.py

from __future__ import annotations

import re
from math import ceil
from typing import Any

from django.core.paginator import EmptyPage, PageNotAnInteger, Paginator
from django.db import models
from django.db.models import CASCADE, PROTECT, SET_NULL, Prefetch
from modelcluster.contrib.taggit import ClusterTaggableManager
from modelcluster.fields import ParentalKey
from taggit.models import Tag, TaggedItemBase
from wagtail.admin.panels import FieldPanel, PageChooserPanel
from wagtail.fields import RichTextField, StreamField
from wagtail.models import Page
from wagtailseo.models import SeoMixin

from apps.blog.blocks import ARTICLE_BODY_BLOCKS


class HomePage(Page):
    featured_article = models.ForeignKey(
        "blog.ArticlePage", null=True, blank=True, on_delete=SET_NULL, related_name="+"
    )

    subpage_types = ["blog.ArticleIndexPage", "legal.LegalIndexPage", "blog.AboutPage"]

    content_panels = Page.content_panels + [
        PageChooserPanel("featured_article", "blog.ArticlePage"),
    ]

    def get_context(self, request, *args, **kwargs):
        ctx = super().get_context(request, *args, **kwargs)
        articles_qs = (
            ArticlePage.objects.live()
            .public()
            .select_related("author")
            .prefetch_related("tags__metadata")
            .order_by("-first_published_at")
        )
        articles = list(articles_qs[:5])
        ctx["featured_article"] = self.featured_article
        ctx["latest_articles"] = articles
        ctx["more_articles"] = articles[:3]
        return ctx


class ArticleIndexPage(Page):
    parent_page_types = ["blog.HomePage"]
    subpage_types = ["blog.ArticlePage"]
    ARTICLES_PER_PAGE = 12

    def get_articles(self):
        return (
            ArticlePage.objects.child_of(self)
            .live()
            .select_related("author")
            .prefetch_related("tags__metadata")
            .order_by("-first_published_at")
        )

    def get_context(self, request, *args, **kwargs):
        ctx = super().get_context(request, *args, **kwargs)
        tag_slug = request.GET.get("tag")
        articles = self.get_articles()
        available_tags = (
            Tag.objects.filter(id__in=articles.values_list("tags__id", flat=True)).distinct().order_by("name")
        )
        if tag_slug:
            articles = articles.filter(tags__slug=tag_slug)
        paginator = Paginator(articles, self.ARTICLES_PER_PAGE)
        page_num = request.GET.get("page")
        try:
            page_obj = paginator.page(page_num)
        except PageNotAnInteger:
            page_obj = paginator.page(1)
        except EmptyPage:
            page_obj = paginator.page(paginator.num_pages)
        ctx["articles"] = page_obj
        ctx["paginator"] = paginator
        ctx["active_tag"] = tag_slug
        ctx["available_tags"] = available_tags
        return ctx


class ArticleTag(TaggedItemBase):
    content_object = ParentalKey("blog.ArticlePage", related_name="tagged_items", on_delete=CASCADE)


class TagMetadata(models.Model):
    COLOUR_CHOICES = [("cyan", "Cyan"), ("pink", "Pink"), ("neutral", "Neutral")]

    tag = models.OneToOneField("taggit.Tag", on_delete=CASCADE, related_name="metadata")
    colour = models.CharField(max_length=20, choices=COLOUR_CHOICES, default="neutral")

    @classmethod
    def get_fallback_css(cls) -> dict[str, str]:
        return {"bg": "bg-zinc-100", "text": "text-zinc-800"}

    def get_css_classes(self) -> dict[str, str]:
        mapping = {
            "cyan": {"bg": "bg-cyan-100", "text": "text-cyan-900"},
            "pink": {"bg": "bg-pink-100", "text": "text-pink-900"},
            "neutral": self.get_fallback_css(),
        }
        return mapping.get(self.colour, self.get_fallback_css())


class ArticlePage(SeoMixin, Page):
    author = models.ForeignKey("authors.Author", on_delete=PROTECT)
    hero_image = models.ForeignKey(
        "wagtailimages.Image", null=True, blank=True, on_delete=SET_NULL, related_name="+"
    )
    summary = models.TextField()
    body = StreamField(ARTICLE_BODY_BLOCKS, use_json_field=True)
    tags = ClusterTaggableManager(through="blog.ArticleTag", blank=True)
    read_time_mins = models.PositiveIntegerField(editable=False, default=1)
    comments_enabled = models.BooleanField(default=True)

    parent_page_types = ["blog.ArticleIndexPage"]
    subpage_types: list[str] = []

    content_panels = Page.content_panels + [
        FieldPanel("author"),
        FieldPanel("hero_image"),
        FieldPanel("summary"),
        FieldPanel("body"),
        FieldPanel("tags"),
        FieldPanel("comments_enabled"),
    ]

    promote_panels = Page.promote_panels + SeoMixin.seo_panels

    search_fields = Page.search_fields

    def save(self, *args: Any, **kwargs: Any) -> None:
        self.read_time_mins = self._compute_read_time()
        return super().save(*args, **kwargs)

    def _compute_read_time(self) -> int:
        words = []
        for block in self.body:
            if block.block_type == "code":
                continue
            value = block.value
            text = value.source if hasattr(value, "source") else str(value)
            words.extend(re.findall(r"\w+", text))
        return max(1, ceil(len(words) / 200))

    def get_tags_with_metadata(self):
        tags = self.tags.all()
        return [(tag, getattr(tag, "metadata", None)) for tag in tags]

    def get_related_articles(self, count: int = 3):
        tag_ids = self.tags.values_list("id", flat=True)
        related = list(
            ArticlePage.objects.live()
            .filter(tags__in=tag_ids)
            .exclude(pk=self.pk)
            .distinct()
            .order_by("-first_published_at")[:count]
        )
        if len(related) < count:
            exclude_ids = [a.pk for a in related] + [self.pk]
            fallback = list(
                ArticlePage.objects.live()
                .exclude(pk__in=exclude_ids)
                .order_by("-first_published_at")[: count - len(related)]
            )
            return related + fallback
        return related

    def get_context(self, request, *args, **kwargs):
        ctx = super().get_context(request, *args, **kwargs)
        ctx["related_articles"] = self.get_related_articles()
        from apps.comments.models import Comment

        approved_replies = Comment.objects.filter(is_approved=True).select_related("parent")
        ctx["approved_comments"] = self.comments.filter(is_approved=True, parent__isnull=True).prefetch_related(
            Prefetch("replies", queryset=approved_replies)
        )
        return ctx


class AboutPage(Page):
    mission_statement = models.TextField()
    body = RichTextField(blank=True)
    featured_author = models.ForeignKey(
        "authors.Author", null=True, blank=True, on_delete=SET_NULL, related_name="+"
    )

    parent_page_types = ["blog.HomePage"]
    subpage_types: list[str] = []

    content_panels = Page.content_panels + [
        FieldPanel("mission_statement"),
        FieldPanel("body"),
        FieldPanel("featured_author"),
    ]