# Source code for debusine.db.models.worker_pools

    # Copyright © The Debusine Developers
    # See the AUTHORS file at the top-level directory of this distribution
    #
    # This file is part of Debusine. It is subject to the license terms
    # in the LICENSE file found in the top-level directory of this
    # distribution. No part of Debusine, including this file, may be copied,
    # modified, propagated, or distributed except according to the terms
    # contained in the LICENSE file.

    """Data models for pools of workers."""

    import logging
    from functools import cached_property
    from typing import TYPE_CHECKING, override

    import pydantic
    from django.core.exceptions import ValidationError
    from django.db import models, transaction
    from django.db.models import JSONField, QuerySet
    from django.urls import reverse
    from django.utils import timezone

    from debusine.db.models.scopes import Scope
    from debusine.server.worker_pools import (
        ScopeWorkerPoolLimits,
        WorkerPoolInterface,
        WorkerPoolLimits,
        WorkerPoolSpecifications,
        provider_interface,
        worker_pool_specifications_model,
    )

    if TYPE_CHECKING:
        from django.http import HttpRequest
        from django_stubs_ext.db.models import TypedModelMeta

        from debusine.db.models.workers import Worker
        from debusine.web.views.ui.worker_pools import WorkerPoolUI
    else:
        TypedModelMeta = object

    logger = logging.getLogger(__name__)

    class WorkerPoolManager(models.Manager["WorkerPool"]):
        """Manager for WorkerPool model."""

        def enabled(self) -> QuerySet["WorkerPool"]:
            """Return connected workers."""
            return WorkerPool.objects.filter(enabled=True)

    class WorkerPool(models.Model):
        """Database model of a worker pool."""

        name = models.SlugField(
            unique=True,
            help_text='Human readable name of the worker pool',
        )
        provider_account = models.ForeignKey("Asset", on_delete=models.PROTECT)
        enabled = models.BooleanField(default=True)
        architectures = JSONField(default=list)
        tags = JSONField(default=list, blank=True)
        specifications = JSONField(default=dict)
        instance_wide = models.BooleanField(default=True)
        ephemeral = models.BooleanField(default=False)
        limits = JSONField(default=dict, blank=True)
        apt_authentication = models.ForeignKey(
            "Asset",
            on_delete=models.PROTECT,
            related_name="+",
            null=True,
            blank=True,
        )

        registered_at = models.DateTimeField()

        objects = WorkerPoolManager()

        class Meta(TypedModelMeta):
            base_manager_name = "objects"

        @override
        def __str__(self) -> str:
            """Return the id and name of the WorkerPool."""
            return f"Id: {self.id} Name: {self.name}"

        def get_absolute_url(self) -> str:
            """Return an absolute URL to view this WorkerPool."""
            return reverse("worker-pools:detail", kwargs={"name": self.name})

        def ui(self, request: "HttpRequest") -> "WorkerPoolUI":
            """Return a UI helper for this instance."""
            from debusine.web.views.ui.worker_pools import WorkerPoolUI

            return WorkerPoolUI.get(request, self)

        @property
        def limits_model(self) -> WorkerPoolLimits:
            """Return the pydantic model for limits."""
            return WorkerPoolLimits.model_validate(self.limits)

        @property
        def specifications_model(self) -> WorkerPoolSpecifications:
            """Return the pydantic model for specifications."""
            model = worker_pool_specifications_model(self.specifications)
            provider_model = self.provider_account.data_model
            assert hasattr(provider_model, "provider_type")
            if provider_model.provider_type != model.provider_type:
                raise ValueError(
                    f"specifications for worker_pool {self.name} do not have a "
                    f"provider_account with a matching provider_type."
                )
            return model

        @property
        def workers_running(self) -> QuerySet["Worker"]:
            """Return the Worker instances that are currently running."""
            return self.worker_set.filter(instance_created_at__isnull=False)

        @property
        def workers_stopped(self) -> QuerySet["Worker"]:
            """Return the Worker instances that are currently stopped."""
            return self.worker_set.filter(instance_created_at__isnull=True)

        @cached_property
        def provider_interface(self) -> WorkerPoolInterface:
            """Return a WorkerPoolInterface instance for this pool."""
            return provider_interface(self)

        def launch_workers(
            self, count: int, override_disabled: bool = False
        ) -> None:
            """Launch count additional worker instances."""
            from debusine.db.models.auth import Token
            from debusine.db.models.workers import Worker

            if not (self.enabled or override_disabled):
                raise ValueError("Pool is disabled, refusing to launch workers.")

            available = self.workers_stopped.count()
            if available < count:
                Worker.objects.create_pool_members(self, count - available)

            launched = 0
            while launched < count:
                with transaction.atomic():
                    worker = (
                        self.workers_stopped.select_for_update(skip_locked=True)
                        .order_by("name")
                        .first()
                    )
                    if worker is None:  # pragma: no cover
                        # Locked or deleted since we expanded the pool, above
                        return

                    old_activation_token = worker.activation_token
                    worker.activation_token = (
                        Token.objects.create_worker_activation()
                    )
                    worker.save()
                    if old_activation_token is not None:
                        old_activation_token.delete()
                    if worker.token is not None:
                        worker.token.disable()

                    self.provider_interface.launch_worker(worker)
                    launched += 1

        def terminate_worker(self, worker: "Worker") -> None:
            """Terminate the specified worker instance."""
            # Import here to prevent circular imports
            from debusine.db.context import context
            from debusine.db.models.work_requests import (
                CannotRetry,
                WorkRequest,
                WorkRequestRetryReason,
            )

            if worker.worker_pool != self:
                raise ValueError(
                    f"pool {self} cannot terminate worker"
                    f" for pool {worker.worker_pool}"
                )

            # Commit this early to avoid scheduling any more work on the worker
            with transaction.atomic():
                if worker.activation_token is not None:
                    worker.activation_token.disable()
                if worker.token is not None:
                    worker.token.disable()

            if worker.instance_created_at is not None:
                WorkerPoolStatistics.objects.create(
                    worker_pool=self,
                    worker=worker,
                    runtime=int(
                        (
                            timezone.now() - worker.instance_created_at
                        ).total_seconds()
                    ),
                )

            # Trigger worker termination
            with transaction.atomic():
                self.provider_interface.terminate_worker(worker)

            with transaction.atomic():
                # De-assign any pending tasks assigned to the worker
                for pending in WorkRequest.objects.reassignable(worker=worker):
                    pending.de_assign_worker()

                # Retry any work requests that were previously running on the worker
                for running in WorkRequest.objects.running(worker=worker):
                    running.mark_aborted()
                    try:
                        with context.disable_permission_checks():
                            running.retry(
                                reason=WorkRequestRetryReason.WORKER_FAILED
                            )
                    except CannotRetry as e:
                        logger.debug(  # noqa: G200
                            "Cannot retry previously-running work request: %s", e
                        )

        @override
        def clean(self) -> None:
            """
            Ensure that data is valid for this worker pool.

            :raise ValidationError: for invalid data.
            """
            try:
                self.limits_model
                self.specifications_model
            except pydantic.ValidationError as e:
                raise ValidationError(message=str(e)) from e

    class ScopeWorkerPool(models.Model):
        """Through table for linking a WorkerPool to a Scope."""

        worker_pool = models.ForeignKey(WorkerPool, on_delete=models.CASCADE)
        scope = models.ForeignKey(Scope, on_delete=models.CASCADE)

        priority = models.IntegerField(default=0)
        limits = JSONField(default=dict, blank=True)

        @override
        def __str__(self) -> str:
            """Return the id and name of the ScopeWorkerPool."""
            return (
                f"Id: {self.id} WorkerPool: {self.worker_pool.name} "
                f"Scope: {self.scope.name}"
            )

        @property
        def limits_model(self) -> ScopeWorkerPoolLimits:
            """Return the pydantic model for limits."""
            return ScopeWorkerPoolLimits.model_validate(self.limits)

    class WorkerPoolTaskExecutionStatistics(models.Model):
        """Time spent executing tasks in a scope, stored at completion."""

        worker_pool = models.ForeignKey(WorkerPool, on_delete=models.CASCADE)
        worker = models.ForeignKey(
            "Worker", null=True, blank=True, on_delete=models.SET_NULL
        )
        scope = models.ForeignKey(Scope, on_delete=models.CASCADE)
        timestamp = models.DateTimeField(auto_now_add=True)
        runtime = models.IntegerField()

        class Meta(TypedModelMeta):
            indexes = [
                models.Index(
                    "timestamp", name="%(app_label)s_worker_pool_exec_ts_idx"
                ),
            ]

    class WorkerPoolStatistics(models.Model):
        """Running time for historical worker instances, stored at shutdown."""

        worker_pool = models.ForeignKey(WorkerPool, on_delete=models.CASCADE)
        worker = models.ForeignKey(
            "Worker", null=True, blank=True, on_delete=models.SET_NULL
        )
        timestamp = models.DateTimeField(auto_now_add=True)
        runtime = models.IntegerField()

        class Meta(TypedModelMeta):
            indexes = [
                models.Index(
                    "timestamp", name="%(app_label)s_worker_pool_stat_ts_idx"
                ),
            ]