import logging
from datetime import timedelta

from celery import shared_task
from django.utils import timezone

from labs.models import LabsOAuthTiDBCloudGrant
from labs.tidb_serverless.oauth import TiDBCloudOAuthClient
from labs.tidb_serverless.exceptions import (
    TiDBCloudUnauthorizedException,
    TiDBCloudRateLimitException,
)

logger = logging.getLogger(__name__)


@shared_task()
def ensure_tidbcloud_oauth_grants_revoked_after_lab():
    # Lab instance will be recycled after 2 hours of creation,
    # thus 3 hours to revoke the oauth grant is enough.
    now = timezone.now()
    created_before = now - timedelta(hours=3)

    for grant in LabsOAuthTiDBCloudGrant.objects.filter(
        revoked_at__isnull=True, created_at__lt=created_before
    ):
        try:
            oauth_client = TiDBCloudOAuthClient(grant)
            oauth_client.revoke_oauth_token()

            try:
                oauth_client.whoami()
                logger.error(
                    f"TiDB Cloud OAuth grant {grant.id} was revoked but token still appears valid."
                )
            except TiDBCloudUnauthorizedException:
                logger.info(f"Successfully revoked TiDB Cloud OAuth grant {grant.id}")
        except TiDBCloudRateLimitException:
            logger.info(
                f"TiDB Cloud API rate limit exceeded, please try running this task again later."
            )
            return
        except Exception as e:
            logger.error(
                f"Failed to revoke TiDB Cloud OAuth grant {grant.id}: {e}",
                exc_info=True,
            )