import uuid
import time

from django.core.management.base import BaseCommand
from django.utils import timezone

from labs.tidb_serverless.api_client import TiDBCloudServerlessClient
from labs.models import (
    ServerlessMainCluster,
    ServerlessBranchCluster,
)


class Command(BaseCommand):
    """
    Ensure a specified main sandbox cluster has the target number of branches.

    This command checks the current number of branches for a given main cluster
    and creates additional branches if needed to meet the target count. It will
    not create branches if the current count already meets or exceeds the target.

    Args:
        --main-cluster-id: The TiDB Serverless main cluster ID to ensure branch count for
        --target-count: The target number of branches that should exist for this cluster
    """

    def add_arguments(self, parser):
        parser.add_argument(
            "--main-cluster-id", type=str, required=True, dest="main_cluster_id"
        )
        parser.add_argument(
            "--target-count", type=int, required=True, dest="target_count"
        )

    def handle(self, *args, **options):
        main_cluster_id = options["main_cluster_id"]
        target_count = options["target_count"]

        try:
            main_cluster = ServerlessMainCluster.objects.get(
                main_cluster_id=main_cluster_id
            )
        except ServerlessMainCluster.DoesNotExist:
            self.stderr.write(f"Main cluster {main_cluster_id} does not exist")
            return

        client = TiDBCloudServerlessClient()

        self.stdout.write(
            f"Start creating new branches for main cluster {main_cluster_id}...\n"
            "NOTE: It will take about 2 minutes to create each branch, please wait patiently..."
        )
        current_count = ServerlessBranchCluster.objects.filter(
            main_cluster_id=main_cluster_id
        ).count()
        created_count = 0
        while current_count < target_count:
            new_branch_name = f"{main_cluster.cluster_name}-{uuid.uuid4().hex}"
            new_branch_id = client.create_branch(main_cluster_id, new_branch_name)
            branch_info = self.wait_for_tidb_serverless_branch_ready(
                client, main_cluster_id, new_branch_id
            )
            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=main_cluster.project_id,
                cluster_id=main_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=new_branch_id,
                main_cluster_id=main_cluster_id,
                cloud_provider=main_cluster.cloud_provider,
                region_name=main_cluster.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(),
            )
            created_count += 1
            current_count = ServerlessBranchCluster.objects.filter(
                main_cluster_id=main_cluster_id
            ).count()
            self.stdout.write(
                f"Created new branch {new_branch_name}! ({current_count}/{target_count})"
            )
        self.stdout.write(f"Done, main cluster {main_cluster_id} has {current_count} branches")

    def wait_for_tidb_serverless_branch_ready(
        self,
        client: TiDBCloudServerlessClient,
        main_cluster_id: str,
        branch_id: str,
        timeout_seconds: int = 600,
    ):
        start_time = time.time()
        end_time = start_time + timeout_seconds
        sleep_seconds = 10
        while time.time() < end_time:
            try:
                branch_info = client.get_branch(main_cluster_id, branch_id)
            except Exception as e:
                self.stderr.write(
                    f"Error getting branch info: {e}, retrying(sleep {sleep_seconds}s)..."
                )
                time.sleep(sleep_seconds)
                continue
            if branch_info["state"] == "ACTIVE":
                return branch_info
            else:
                time.sleep(sleep_seconds)
        raise Exception(
            f"Branch ({branch_id}) cannot be ready after {timeout_seconds} seconds"
        )