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

"""Middleware for updating Worker.last_seen_at."""

from collections.abc import Awaitable, Callable, Iterable
from typing import TYPE_CHECKING

from asgiref.typing import ASGIReceiveCallable, ASGISendCallable
from channels.db import database_sync_to_async
from django.utils import timezone

from debusine.db.context import context
from debusine.db.models import Token, WorkRequest
from debusine.db.permissioncontext import PermissionContext
from debusine.server.consumers import ChannelScopeWithToken

if TYPE_CHECKING:
    from channels.consumer import _ChannelScope


class TokenLastSeenAtMiddlewareChannels:
    """
    Save the Token in the scope and update Worker.last_seen_at.

    Token available in the scope helps to avoid double lookups.
    """

    def __init__(
        self,
        app: Callable[
            [ChannelScopeWithToken, ASGIReceiveCallable, ASGISendCallable],
            Awaitable[None],
        ],
    ) -> None:
        """Middleware API entry point."""
        self.app = app

    @staticmethod
    def get_token_header(headers: Iterable[tuple[bytes, bytes]]) -> str | None:
        """Return the token key from the headers or None."""
        for header in headers:
            if header[0].lower() == b"token":
                return header[1].decode("utf-8")

        return None

    def init_context_with_work_request_token(self, token: Token) -> None:
        """Initialize the context from a work request token."""
        # This is a sync operation, as it will do a database lookup to precache
        # scope and workspace roles
        assert token.work_request is not None
        context.set_permission_context_scoped(
            token.work_request.workspace.scope,
            PermissionContext.for_work_request(token=token),
        )

    async def __call__(
        self,
        scope: "_ChannelScope",
        receive: ASGIReceiveCallable,
        send: ASGISendCallable,
    ) -> None:
        """Middleware entry point."""
        # When middleware is modifying the scope, it should make a copy of the
        # scope object before mutating it.
        # See https://asgi.readthedocs.io/en/latest/specs/main.html#middleware
        scope_with_token = ChannelScopeWithToken(**scope, token=None)

        token_key = self.get_token_header(scope_with_token["headers"])

        if token_key is None:
            # No token was in the request, nothing needs to be done
            scope_with_token["token"] = None
            return await self.app(scope_with_token, receive, send)

        token = await database_sync_to_async(
            Token.objects.select_related(
                "work_request__workspace__scope",
                "work_request__created_by",
            ).get_token_or_none
        )(token_key)

        if token is not None:
            token.last_seen_at = timezone.now()
            await token.asave()
            if token.enabled:
                match token.token_type:
                    case Token.Types.WORKER:
                        context.set_permission_context_scopeless(
                            PermissionContext.for_worker(token)
                        )
                    case Token.Types.WORK_REQUEST:
                        if (
                            token.work_request is not None
                            and token.work_request.status
                            == WorkRequest.Statuses.RUNNING
                        ):
                            await database_sync_to_async(
                                self.init_context_with_work_request_token
                            )(token)
                    case _:
                        pass

        scope_with_token["token"] = token

        return await self.app(scope_with_token, receive, send)
