from django.core.management.base import BaseCommand

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


class Command(BaseCommand):
    """
    Clean up orphaned serverless sandbox branches by synchronizing database records with TiDB Cloud.

    This command performs a bidirectional cleanup:
    1. Removes orphaned branches in TiDB Cloud that have no corresponding database records
    2. Removes stale database records for branches that no longer exist in TiDB Cloud

    This helps maintain data consistency between the local database and TiDB Cloud,
    preventing resource leaks and ensuring accurate branch tracking.

    Usage:
        python manage.py cleanup_orphaned_sandbox_branches --main-cluster-id <cluster_id>

    Args:
        --main-cluster-id: The TiDB Serverless main cluster ID to clean up branches for
    """

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

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

        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

        db_branch_ids = set(
            ServerlessBranchCluster.objects.filter(
                main_cluster_id=main_cluster_id
            ).values_list("branch_cluster_id", flat=True)
        )
        cloud_branch_ids = set()
        page_token = ""
        while True:
            result = client.list_branches(main_cluster_id, page_token=page_token)
            cloud_branch_ids.update(
                [branch["branchId"] for branch in result["branches"]]
            )
            page_token = result["nextPageToken"]
            if not page_token:
                break

        not_in_db = cloud_branch_ids - db_branch_ids
        not_in_cloud = db_branch_ids - cloud_branch_ids

        self.stdout.write(
            f"There are {len(cloud_branch_ids)} branches in cloud, {len(db_branch_ids)} branches in db"
        )
        if not_in_db:
            self.stdout.write(
                "\n"
                + "=" * 10
                + f" Found {len(not_in_db)} branches in cloud but not in db "
                + "=" * 10
                + "\n"
            )
            for branch_id in not_in_db:
                try:
                    client.delete_branch(main_cluster_id, branch_id)
                    self.stdout.write(f"Deleted branch {branch_id} in cloud!")
                except Exception as e:
                    self.stdout.write(
                        f"Failed to delete branch {branch_id} in cloud: {e}"
                    )

        if not_in_cloud:
            self.stdout.write(
                "\n"
                + "=" * 10
                + f" Found {len(not_in_cloud)} branches in db but not in cloud "
                + "=" * 10
                + "\n"
            )
            ServerlessBranchCluster.objects.filter(
                branch_cluster_id__in=not_in_cloud
            ).delete()
            self.stdout.write(
                f"Deleted {len(not_in_cloud)} branches in db but not in cloud!"
            )