# 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.

"""
Backend infrastructure for provisioning cloud workers.

See docs/reference/devel-blueprints/dynamic-compute.rst
"""

import logging
import math
from collections import defaultdict
from functools import cached_property
from typing import NamedTuple

from celery import shared_task
from django.contrib.postgres.functions import TransactionNow
from django.db import transaction
from django.db.models import Count, F, IntegerField, Min, Q, Sum
from django.db.models.functions import Cast, Extract, Greatest, Least
from django.utils import timezone

from debusine.db.context import context
from debusine.db.models import (
    Scope,
    ScopeWorkerPool,
    WorkRequest,
    Worker,
    WorkerPool,
)
from debusine.db.models.worker_pools import (
    WorkerPoolStatistics,
    WorkerPoolTaskExecutionStatistics,
)
from debusine.server.scheduler import (
    get_candidate_work_requests,
    get_provided_worker_tags,
    get_required_task_tags,
)
from debusine.tasks.models import WorkerType

# TODO: Hardcoded estimation of duration for a task, until we have code to
#       estimate the duration of future tasks based on task-statistics
TASK_DURATION_SECONDS = 10 * 60

# TODO: Hardcoded target latency for a scope, until we can read it from Scope
TARGET_LATENCY_SECONDS = 3600

log = logging.getLogger(__name__)


class TagSetPair(NamedTuple):
    """The two sets of tags used for scheduling calculations."""

    required: frozenset[str]
    provided: frozenset[str]


def terminate_idle(minutes: int = 10) -> bool:
    """
    Find and terminate idle workers.

    :returns: true if idle workers were found
    """
    # List cloud workers that have been idle for more than 10 minutes
    idle_workers = (
        Worker.objects.active()
        .with_idle_time()
        .filter(
            idle_time__gt=minutes * 60,
            worker_pool__isnull=False,
            worker_pool__enabled=True,
        )
    )
    has_idle = False
    for worker in idle_workers:
        assert worker.worker_pool is not None
        has_idle = True
        log.info(
            "Terminating idle worker %s (idle for %d seconds)",
            worker.name,
            worker.idle_time,
        )
        worker.worker_pool.terminate_worker(worker)
    return has_idle


class WorkerAllocator:
    """Allocate new workers needed to handle the pending load."""

    # NOTE: This is a minimum viable implementation of cloud compute
    # provisioning.
    #
    # The following simplifications are in place:
    # * assume a hardcoded running time for tasks, until we have code to
    #   estimate the duration of future tasks based on task-statistics
    # * assume all workers take the same time to process a task

    def __init__(self) -> None:
        """Prefetch common information from the database."""
        self.start_of_month = timezone.now().replace(
            day=1, hour=0, minute=0, second=0, microsecond=0
        )
        self.static_workers = self.static_workers_count()

    def static_workers_count(self) -> dict[TagSetPair, int]:
        """Count connected static workers by provided_tags."""
        workers: dict[TagSetPair, int] = defaultdict(int)
        aggregates = list(
            Worker.objects.connected()
            .filter(worker_pool__isnull=True, worker_type=WorkerType.EXTERNAL)
            .with_scheduler_tags()
            .values("provided_tags", "required_tags")
            .annotate(count=Count("id"), example_id=Min("id"))
            .order_by()
        )

        examples: dict[int, Worker] = {}
        for worker in Worker.objects.filter(
            id__in=[row["example_id"] for row in aggregates]
        ):
            examples[worker.id] = worker

        for row in aggregates:
            example = examples[row["example_id"]]
            tagset = TagSetPair(
                required=frozenset(get_required_task_tags(example)),
                provided=frozenset(get_provided_worker_tags(example)),
            )
            workers[tagset] = row["count"]
        return workers

    def limit_by_max_active_instances(
        self, pool: WorkerPool, amount: int
    ) -> int:
        """Limit amount by the max_active_instances configured in the pool."""
        pool_limits = pool.limits_model
        if pool_limits.max_active_instances is None:
            return amount

        current_worker_count = pool.workers_running.count()
        available = pool_limits.max_active_instances - current_worker_count
        if available <= 0:
            return 0
        return min(amount, available)

    def get_used_seconds(self, pool: WorkerPool) -> int:
        """Count the seconds of pool usage for this month."""
        completed = (
            WorkerPoolStatistics.objects.filter(
                worker_pool=pool, timestamp__gte=self.start_of_month
            )
            .annotate(
                runtime_this_month=Least(
                    "runtime",
                    Cast(
                        Extract(F("timestamp") - self.start_of_month, "epoch"),
                        IntegerField(),
                    ),
                )
            )
            .aggregate(seconds_per_month=Sum("runtime_this_month"))[
                "seconds_per_month"
            ]
            or 0
        )
        connected = (
            Worker.objects.connected()
            .filter(worker_pool=pool)
            .annotate(
                runtime=Extract(
                    TransactionNow()
                    - Greatest("instance_created_at", self.start_of_month),
                    "epoch",
                )
            )
            .aggregate(seconds_per_month=Sum("runtime"))["seconds_per_month"]
            or 0
        )
        return completed + connected

    def limit_by_pool_target_max_seconds_per_month(
        self, pool: WorkerPool, amount: int
    ) -> int:
        """Limit amount by the pool's target_max_seconds_per_month."""
        target_max_seconds_per_month: int | None = (
            pool.limits_model.target_max_seconds_per_month
        )
        if target_max_seconds_per_month is None:
            return amount

        used_seconds = self.get_used_seconds(pool)
        if used_seconds >= target_max_seconds_per_month:
            return 0

        return min(
            amount,
            int(
                (target_max_seconds_per_month - used_seconds)
                / TASK_DURATION_SECONDS
            ),
        )

    def process_scope(self, scope: Scope) -> None:
        """Perform allocation for a scope."""
        allocator = ScopeWorkerAllocator(self, scope)
        allocator.process_scope()


class ScopeWorkerAllocator:
    """Allocate new workers needed to handle the pending load of a scope."""

    def __init__(self, worker_allocator: WorkerAllocator, scope: Scope) -> None:
        """Prefetch common information from the database."""
        self.worker_allocator = worker_allocator
        self.scope = scope

    def get_used_seconds(self, pool: WorkerPool) -> int:
        """Count the seconds of pool usage for the current scope this month."""
        completed = (
            WorkerPoolTaskExecutionStatistics.objects.filter(
                worker_pool=pool,
                scope=self.scope,
                timestamp__gte=self.worker_allocator.start_of_month,
            )
            .annotate(
                runtime_this_month=Least(
                    "runtime",
                    Cast(
                        Extract(
                            F("timestamp")
                            - self.worker_allocator.start_of_month,
                            "epoch",
                        ),
                        IntegerField(),
                    ),
                )
            )
            .aggregate(seconds_per_month=Sum("runtime_this_month"))[
                "seconds_per_month"
            ]
            or 0
        )
        running = (
            WorkRequest.objects.running()
            .filter(
                worker__worker_pool=pool,
                workspace__scope=self.scope,
                started_at__isnull=False,
            )
            .annotate(
                runtime=Extract(
                    TransactionNow()
                    - Greatest(
                        "started_at", self.worker_allocator.start_of_month
                    ),
                    "epoch",
                )
            )
            .aggregate(seconds_per_month=Sum("runtime"))["seconds_per_month"]
            or 0
        )
        return completed + running

    def limit_by_scope_pool_target_max_seconds_per_month(
        self, pool: WorkerPool, amount: int
    ) -> int:
        """Limit amount by target_max_seconds_per_month for scope and pool."""
        target_max_seconds_per_month: int | None = None
        try:
            scope_worker_pool = ScopeWorkerPool.objects.get(
                worker_pool=pool, scope=self.scope
            )
        except ScopeWorkerPool.DoesNotExist:
            return amount

        target_max_seconds_per_month = (
            scope_worker_pool.limits_model.target_max_seconds_per_month
        )
        if target_max_seconds_per_month is None:
            return amount

        used_seconds = self.get_used_seconds(pool)
        if used_seconds >= target_max_seconds_per_month:
            return 0

        return min(
            amount,
            int(
                (target_max_seconds_per_month - used_seconds)
                / TASK_DURATION_SECONDS
            ),
        )

    def allocate_workers(self, pool: WorkerPool, amount: int) -> int:
        """
        Allocate workers on pool for the given scope.

        The number of workers allocated can be less than amount, or even zero,
        depending on pool availability and quotas.

        :param scope: the scope requesting allocation
        :param pool: the pool used to allocate
        :param amount: the amount of new workers required
        :returns: the amount of workers allocated
        """
        # Enforce max_active_instances
        amount = self.worker_allocator.limit_by_max_active_instances(
            pool, amount
        )

        # Enforce target_max_seconds_per_month
        amount = (
            self.worker_allocator.limit_by_pool_target_max_seconds_per_month(
                pool, amount
            )
        )

        # Enforce target_max_seconds_per_month
        amount = self.limit_by_scope_pool_target_max_seconds_per_month(
            pool, amount
        )

        if amount > 0:
            log.info("Launching %d workers in %s", amount, pool.name)
            pool.launch_workers(amount)
        return amount

    def workers_required_for_pool(self, pool: WorkerPool) -> int:
        """
        Compute how many workers from pool can process tasks for scope.

        This computes the absolute number of workers required to process the
        scope's pending load for the given worker pool, within the expected
        latency targets.
        """
        # If we have no pending tasks, we do not require workers
        tasks = get_candidate_work_requests(pool).count()
        if tasks == 0:
            return 0

        # Estimated number of worker-seconds needed to run all but the last
        # pending task.
        pending_seconds = TASK_DURATION_SECONDS * max(tasks - 1, 0)

        # Each worker will contribute one second per second of work towards
        # draining the queue.  There's no point in allocating more workers
        # than the number of pending tasks.
        return max(
            1,
            min(math.ceil(pending_seconds / TARGET_LATENCY_SECONDS), tasks),
        )

    def workers_available_for_tagset(self, tagset: TagSetPair) -> int:
        """
        Compute the number of workers with tagset available to the scope.

        Count workers that are able to execute tasks that a worker with tagset
        can execute.
        """
        workers = 0
        for static_ts, count in self.worker_allocator.static_workers.items():
            if tagset.provided.issubset(
                static_ts.provided
            ) and static_ts.required.issubset(tagset.required):
                workers += count

        workers += (
            Worker.objects.active()
            .filter(
                Q(worker_pool__instance_wide=True)
                | Q(worker_pool__scopeworkerpool__scope=self.scope)
            )
            .with_scheduler_tags()
            # mypy_django_plugin doesn't know that this is interpreted using
            # JSON lookups at run-time.
            .filter(provided_tags__contains=list(tagset.provided))  # type: ignore[misc]  # noqa: E501
            .filter(required_tags__contained_by=list(tagset.required))
            .values("id")
        ).count()

        return workers

    @cached_property
    def allocation_pools(self) -> list[WorkerPool]:
        """
        Return the list of pools that can be used to allocate workers.

        The list is ordered by priority, so that allocation can proceed from
        the first element onwards.
        """
        # Pools in which we can allocate them
        pools: list[WorkerPool] = []

        # Add scope pools by decreasing priority
        pools.extend(
            WorkerPool.objects.enabled()
            .filter(scopeworkerpool__scope=self.scope)
            .order_by("-scopeworkerpool__priority")
        )

        # Add instance-wide pools at the end
        pools.extend(WorkerPool.objects.enabled().filter(instance_wide=True))

        # TODO: enforce that ScopeWorkerPool cannot link to instance-wide
        # pools? Otherwise we may need to remove duplicates here

        return pools

    def process_scope(self) -> None:
        """Perform allocation for this scope."""
        # Number of workers required to process the queue in time

        # Allocate from one pool at a time
        for pool in self.allocation_pools:
            workers_required = self.workers_required_for_pool(pool)
            if workers_required == 0:
                continue

            tagset = TagSetPair(
                provided=frozenset(get_provided_worker_tags(pool)),
                required=frozenset(get_required_task_tags(pool)),
            )
            workers_available = self.workers_available_for_tagset(tagset)

            # Number of workers that we need to allocate
            new_workers_required = workers_required - workers_available

            # If we have enough workers already, we are done
            if new_workers_required <= 0:
                continue

            self.allocate_workers(pool, new_workers_required)


def create_new_workers() -> None:
    """Provision new workers if needed to handle pending load."""
    # NOTE: This is a minimum viable implementation of cloud compute
    # provisioning, intended to be useful until we implement tagged scheduling
    # (#326).
    #
    # The following simplifications are in place:
    # * assume scopes' load demands are isolated from each other, which allows
    #   to process each scope's work queue without taking other scopes into
    #   account.
    #   This is not correct with regards to static workers, instance-wide pools
    #   or pools shared by multiple scopes, but in the worst case we
    #   underestimate the number of new workers needed and end up provisioning
    #   some workers again later as queues drain more slowly than estimated.
    #
    #   i.e. if two scopes would each need 15 workers to meet targets and there
    #   are 10 static workers, they'd each provision 5 from pools in this MVP,
    #   while in reality they'd need more because they're both sharing the same
    #   static workers.
    # * see WorkerAllocator for more assumptions

    allocator = WorkerAllocator()
    for scope in Scope.objects.all():
        allocator.process_scope(scope)


def scrub_workers() -> None:
    """Scrub all WorkerPools."""
    for pool in WorkerPool.objects.all():
        pool.provider_interface.terminate_missing()


# mypy complains that celery.shared_task is untyped, which is true, but we
# can't fix that here.
@shared_task()  # type: ignore[untyped-decorator]
@transaction.atomic
@context.local()
@context.disable_permission_checks()
def provision() -> None:
    """Provision or terminate workers from pools to handle pending tasks."""
    # NOTE: This is a minimum viable implementation of cloud compute
    # provisioning, intended to be useful until we implement tagged scheduling
    # (#326).
    #
    # The following simplifications are in place:
    #  * do not consider per-scope quotas when terminating workers: only idle
    #    workers are terminated
    #  * see create_new_workers() comments for more assumptions

    if terminate_idle():
        # If workers have been idle for more than 10 minutes, we can assume
        # that we do not have a significant backlog of pending work requests.
        #
        # This may not be true: we may have a pool of idle arm64 workers and a
        # backlog of amd64 tasks. If that is the case, we will catch up with
        # the backlog in the next provision run.
        return

    create_new_workers()


# mypy complains that celery.shared_task is untyped, which is true, but we
# can't fix that here.
@shared_task()  # type: ignore[untyped-decorator]
@transaction.atomic
@context.local()
@context.disable_permission_checks()
def scrub() -> None:
    """
    Scan the worker pool for failed workers and clean them up.

    We expect some number of workers to die, especially if we're using cheap
    "spot-priced" instances. Handle their death.
    """
    scrub_workers()
