import logging

from labs.lab_manager_v2.tasks.base import BaseLabTask
from labs.tidb_serverless.branch import (
    delete_used_serverless_branch,
    create_new_serverless_branch,
)
from labs.tidb_serverless.oauth import TiDBCloudOAuthClient

logger = logging.getLogger(__name__)


class TiDBCloudServerlessClusterDeleteTask(BaseLabTask):
    def run(self):
        db_lab_instance = self.db_lab_instance
        if db_lab_instance.user_serverless_cluster:
            user_serverless_cluster = db_lab_instance.user_serverless_cluster
            oauth_client = TiDBCloudOAuthClient(user_serverless_cluster.oauth_grant)
            oauth_client.cleanup_lab_cluster(user_serverless_cluster)
        else:
            delete_used_serverless_branch(db_lab_instance.unique_tag)
            create_new_serverless_branch(
                db_lab_instance.region_name,
                db_lab_instance.cloud_provider,
            )