Tussilago/control_plane/admin.py

from __future__ import annotations

from typing import TYPE_CHECKING

from django.contrib import admin
from django.contrib import messages
from django.db.models import Count
from django.db.models import F

from control_plane.models import Deployment
from control_plane.models import HostedSite
from control_plane.models import RuntimeService
from control_plane.models import RuntimeServiceKind
from control_plane.models import Tenant
from control_plane.tasks import provision_test_runtime_services

if TYPE_CHECKING:
    from django.db.models import QuerySet
    from django.http import HttpRequest

    RuntimeServiceInlineBase = admin.StackedInline[RuntimeService]
    TenantAdminBase = admin.ModelAdmin[Tenant]
    HostedSiteAdminBase = admin.ModelAdmin[HostedSite]
    DeploymentAdminBase = admin.ModelAdmin[Deployment]
    RuntimeServiceAdminBase = admin.ModelAdmin[RuntimeService]
else:
    RuntimeServiceInlineBase = admin.StackedInline
    TenantAdminBase = admin.ModelAdmin
    HostedSiteAdminBase = admin.ModelAdmin
    DeploymentAdminBase = admin.ModelAdmin
    RuntimeServiceAdminBase = admin.ModelAdmin


class RuntimeServiceInline(RuntimeServiceInlineBase):
    """Allow deployment admins to create/edit related runtime services inline."""

    model = RuntimeService
    extra = 0
    max_num = len(RuntimeServiceKind)
    show_change_link = True


@admin.register(Tenant)
class TenantAdmin(TenantAdminBase):
    """Expose tenants for admin-managed smoke data setup."""

    list_display = ("slug", "display_name")
    search_fields = ("slug", "display_name")
    ordering = ("slug",)


@admin.register(HostedSite)
class HostedSiteAdmin(HostedSiteAdminBase):
    """Expose hosted sites so admins can build deployment test graphs."""

    list_display = ("slug", "display_name", "tenant_slug", "service_port")
    list_filter = ("tenant",)
    search_fields = (
        "slug",
        "display_name",
        "tenant__slug",
        "tenant__display_name",
        "wsgi_module",
    )
    ordering = ("tenant__slug", "slug")
    autocomplete_fields = ("tenant",)
    list_select_related = ("tenant",)

    def get_queryset(self, request: HttpRequest) -> QuerySet[HostedSite]:
        """Load tenant slug values for changelist rendering.

        Returns:
            Hosted site queryset with tenant join and tenant slug annotation.
        """
        return (
            super()
            .get_queryset(request)
            .select_related("tenant")
            .annotate(
                tenant_slug_value=F("tenant__slug"),
            )
        )

    @admin.display(ordering="tenant__slug", description="Tenant")
    def tenant_slug(self, hosted_site: HostedSite) -> str:
        """Return tenant slug for changelist display and sorting."""
        return str(vars(hosted_site)["tenant_slug_value"])


@admin.register(Deployment)
class DeploymentAdmin(DeploymentAdminBase):
    """Expose deployments and queue test container provisioning."""

    list_display = (
        "id",
        "status",
        "tenant_slug",
        "site_slug",
        "idempotency_key",
        "guest_port",
        "runtime_service_total",
    )
    list_filter = ("status",)
    search_fields = (
        "=id",
        "idempotency_key",
        "firecracker_vm_id",
        "hosted_site__slug",
        "hosted_site__tenant__slug",
    )
    ordering = ("hosted_site__tenant__slug", "hosted_site__slug", "-created_at")
    autocomplete_fields = ("hosted_site",)
    list_select_related = ("hosted_site__tenant",)
    inlines = (RuntimeServiceInline,)
    actions = ("create_test_containers",)

    def get_queryset(self, request: HttpRequest) -> QuerySet[Deployment]:
        """Load related hosted site and tenant rows for admin rendering.

        Returns:
            Deployment queryset with hosted site and tenant joined.
        """
        return (
            super()
            .get_queryset(request)
            .select_related("hosted_site__tenant")
            .annotate(
                tenant_slug_value=F("hosted_site__tenant__slug"),
                site_slug_value=F("hosted_site__slug"),
                runtime_service_total_value=Count("runtime_services", distinct=True),
            )
        )

    @admin.display(ordering="hosted_site__tenant__slug", description="Tenant")
    def tenant_slug(self, deployment: Deployment) -> str:
        """Return tenant slug for changelist display and sorting."""
        return str(vars(deployment)["tenant_slug_value"])

    @admin.display(ordering="hosted_site__slug", description="Site")
    def site_slug(self, deployment: Deployment) -> str:
        """Return hosted site slug for changelist display and sorting."""
        return str(vars(deployment)["site_slug_value"])

    @admin.display(description="Runtime services")
    def runtime_service_total(self, deployment: Deployment) -> int:
        """Return total runtime services currently linked to a deployment."""
        return int(vars(deployment)["runtime_service_total_value"])

    @admin.action(description="Queue test container provisioning")
    def create_test_containers(
        self,
        request: HttpRequest,
        queryset: QuerySet[Deployment],
    ) -> None:
        """Queue Celery jobs that seed and provision local test containers."""
        deployment_ids = [str(deployment_id) for deployment_id in queryset.values_list("id", flat=True)]
        for deployment_id in deployment_ids:
            provision_test_runtime_services.delay(deployment_id)

        self.message_user(
            request,
            (
                f"Queued test container provisioning for {len(deployment_ids)} deployments. "
                "Run a Celery worker to execute queued jobs."
            ),
            level=messages.SUCCESS,
        )


@admin.register(RuntimeService)
class RuntimeServiceAdmin(RuntimeServiceAdminBase):
    """Expose runtime service containers to Django admin users."""

    list_display = (
        "container_name",
        "kind",
        "status",
        "tenant_slug",
        "site_slug",
        "internal_port",
    )
    list_filter = ("kind", "status")
    search_fields = (
        "container_name",
        "network_name",
        "hostname",
        "deployment__idempotency_key",
        "deployment__hosted_site__slug",
        "deployment__hosted_site__tenant__slug",
    )
    ordering = (
        "deployment__hosted_site__tenant__slug",
        "deployment__hosted_site__slug",
        "kind",
    )
    autocomplete_fields = ("deployment",)
    list_select_related = ("deployment__hosted_site__tenant",)

    def get_queryset(self, request: HttpRequest) -> QuerySet[RuntimeService]:
        """Load related deployment context for changelist rendering.

        Returns:
            Runtime service queryset with deployment, site, and tenant joined.
        """
        return (
            super()
            .get_queryset(request)
            .select_related("deployment__hosted_site__tenant")
            .annotate(
                tenant_slug_value=F("deployment__hosted_site__tenant__slug"),
                site_slug_value=F("deployment__hosted_site__slug"),
            )
        )

    @admin.display(ordering="deployment__hosted_site__tenant__slug", description="Tenant")
    def tenant_slug(self, runtime_service: RuntimeService) -> str:
        """Return tenant slug for changelist display and sorting."""
        return str(vars(runtime_service)["tenant_slug_value"])

    @admin.display(ordering="deployment__hosted_site__slug", description="Site")
    def site_slug(self, runtime_service: RuntimeService) -> str:
        """Return hosted site slug for changelist display and sorting."""
        return str(vars(runtime_service)["site_slug_value"])