import logging

from celery import shared_task
from django.utils import timezone

from labs.models import ServerlessBranchCluster

logger = logging.getLogger(__name__)


@shared_task(bind=True, max_retries=20)
def wait_for_tidb_serverless_cluster_ready_and_save(
    self,
    project_id: str,
    cluster_id: str,
    branch_id: str,
    cloud_provider: str,
    region_name: str,
):
    from labs.tidb_serverless.api_client import TiDBCloudServerlessClient

    client = TiDBCloudServerlessClient()
    try:
        branch_info = client.get_branch(cluster_id, branch_id)
    except Exception as e:
        raise Exception(
            f"Failed to get branch info, cluster_id: {cluster_id}, "
            f"branch_id: {branch_id}, region_name: {region_name}, error: {e}"
        ) from e

    if not branch_info:
        raise Exception(
            f"Branch info not found, cluster_id: {cluster_id}, "
            f"branch_id: {branch_id}, region_name: {region_name}"
        )

    if branch_info["state"] != "ACTIVE":
        logger.debug(
            f"Branch is not ready, retry, cluster_id: {cluster_id}, "
            f"branch_id: {branch_id}, region_name: {region_name}"
        )
        # Wait for 60 seconds before retrying
        # TiDB Serverless branch is not ready immediately after creation,
        # sometimes it takes 1-3 minutes to be ready
        raise self.retry(countdown=60)

    branch_name = branch_info["displayName"]
    cluster_host = branch_info["endpoints"]["public"]["host"]
    cluster_port = branch_info["endpoints"]["public"]["port"]

    admin_user_info = client.create_admin_user_for_branch(
        project_id=project_id,
        cluster_id=cluster_id,
        branch_name=branch_name,
    )
    admin_user_name = admin_user_info["username"]
    admin_user_password = admin_user_info["password"]

    ServerlessBranchCluster.objects.create(
        branch_cluster_id=branch_id,
        main_cluster_id=cluster_id,
        cloud_provider=cloud_provider,
        region_name=region_name,
        cluster_name=branch_name,
        cluster_host=cluster_host,
        cluster_port=cluster_port,
        admin_name=admin_user_name,
        admin_password=admin_user_password,
        in_use=ServerlessBranchCluster.InUseChoices.N,
        create_time=timezone.now(),
    )