import uuid
import logging
from typing import Optional

from django.db import transaction
from django.utils import timezone

from labs.models import ServerlessBranchCluster, ServerlessMainCluster
from labs.tidb_serverless.api_client import TiDBCloudServerlessClient
from labs.tidb_serverless.schemas import TiDBServerlessClusterConfig


logger = logging.getLogger(__name__)


def checkout_serverless_branch(
    lab_instance_unique_tag: str, region_name: str, cloud_provider: str = "aws"
) -> Optional[TiDBServerlessClusterConfig]:
    with transaction.atomic():
        available_branch = (
            ServerlessBranchCluster.objects.select_for_update()
            .filter(
                in_use=ServerlessBranchCluster.InUseChoices.N,
                region_name=region_name,
                cloud_provider=cloud_provider,
            )
            .exclude(
                branch_cluster_id="N/A",
                admin_name="N/A",
            )
            .first()
        )
        if not available_branch:
            return None

        # TODO: I don't know why we need to check if the cluster_name is N/A, need to figure out
        if available_branch.cluster_name == "N/A":
            return None

        available_branch.in_use = ServerlessBranchCluster.InUseChoices.Y
        available_branch.trainer_tag = lab_instance_unique_tag
        available_branch.check_out_time = timezone.now()
        available_branch.save(update_fields=["in_use", "trainer_tag", "check_out_time"])
        return TiDBServerlessClusterConfig(
            environment="sandbox",
            cluster_id=available_branch.branch_cluster_id,
            cluster_name=available_branch.cluster_name,
            cluster_host=available_branch.cluster_host,
            cluster_port=available_branch.cluster_port,
            admin_name=available_branch.admin_name,
            admin_password=available_branch.admin_password,
        )


def delete_used_serverless_branch(lab_instance_unique_tag: str):
    try:
        branch = ServerlessBranchCluster.objects.get(
            trainer_tag=lab_instance_unique_tag,
            in_use=ServerlessBranchCluster.InUseChoices.Y,
        )
    except ServerlessBranchCluster.DoesNotExist:
        logger.error(
            f"Serverless branch not found for lab instance unique tag: {lab_instance_unique_tag}"
        )
        return

    client = TiDBCloudServerlessClient()
    client.delete_branch(
        cluster_id=branch.main_cluster_id,
        branch_id=branch.branch_cluster_id,
    )
    branch.delete()


def get_serverless_main_cluster(
    region_name: str, cloud_provider: str = "aws"
) -> Optional[ServerlessMainCluster]:
    return ServerlessMainCluster.objects.filter(
        region_name=region_name, cloud_provider=cloud_provider
    ).first()


def create_new_serverless_branch(
    region_name: str, cloud_provider: str = "aws"
) -> tuple[str, str]:
    main_cluster = get_serverless_main_cluster(region_name, cloud_provider)
    if not main_cluster:
        raise Exception(f"Serverless main cluster not found for region: {region_name}")

    new_branch_name = f"{main_cluster.cluster_name}-{uuid.uuid4().hex}"

    client = TiDBCloudServerlessClient()
    new_branch_id = client.create_branch(
        cluster_id=main_cluster.main_cluster_id,
        new_branch_name=new_branch_name,
    )
    from labs.tasks import tidb_serverless as tidb_serverless_tasks

    tidb_serverless_tasks.wait_for_tidb_serverless_cluster_ready_and_save.apply_async(
        args=(
            main_cluster.project_id,
            main_cluster.main_cluster_id,
            new_branch_id,
            cloud_provider,
            region_name,
        ),
        countdown=30,
    )
    return new_branch_id, new_branch_name