from rest_framework.exceptions import NotFound

from labs.api.base import Endpoint
from labs.labs_core.stateless_functions import get_lab_metadata

from labs.models import (
    LabsOAuthTiDBCloudGrant,
)


class LabDetailBaseEndpoint(Endpoint):
    def convert_args(self, request, lab_code: str, *args, **kwargs):
        try:
            lab_metadata = get_lab_metadata(lab_code)
        except KeyError:
            raise NotFound(detail="Lab not found")

        kwargs["lab_metadata"] = lab_metadata
        kwargs["lab_code"] = lab_code
        return args, kwargs


class LabTiDBCloudGrantBaseEndpoint(Endpoint):
    def convert_args(self, request, grant_id: int, *args, **kwargs):
        try:
            grant = LabsOAuthTiDBCloudGrant.objects.get(
                id=grant_id,
                user=request.user,
            )
        except LabsOAuthTiDBCloudGrant.DoesNotExist:
            raise NotFound(detail="Grant not found")

        kwargs["grant"] = grant
        return args, kwargs