import MySQLdb
from django.core.management.base import BaseCommand
from tqdm import tqdm

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


class Command(BaseCommand):
    def add_arguments(self, parser):
        parser.add_argument(
            "--main-cluster-id", type=str, required=True, dest="main_cluster_id"
        )
        parser.add_argument(
            "--cert-path",
            type=str,
            required=False,
            dest="cert_path",
            default="/etc/ssl/certs/ca-certificates.crt",
        )

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

        try:
            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

        for branch in tqdm(
            ServerlessBranchCluster.objects.filter(
                main_cluster_id=main_cluster_id
            ).iterator(chunk_size=20),
            desc="Checking",
            unit="branches",
        ):
            try:
                db = MySQLdb.connect(
                    host=branch.cluster_host,
                    port=int(branch.cluster_port),
                    user=branch.admin_name,
                    password=branch.admin_password,
                    database="test",
                    ssl_mode="VERIFY_IDENTITY",
                    ssl={
                        "ca": cert_path,
                    },
                )
                cursor = db.cursor()
                cursor.execute("SELECT 1")
                result = cursor.fetchone()
                assert result == (1,)
            except Exception as e:
                self.stderr.write(f"Branch {branch.cluster_name} is not available: {e}")