from time import sleep
import boto3
from tqdm import tqdm
from google.cloud import bigtable
from google.cloud.bigtable import enums


def cleanup_dynamo_ci():
    db = boto3.resource("dynamodb")

    num_to_delete = 0
    all_tables = db.tables.all()
    for table in all_tables:
        if "integration_test" in table.name:
            num_to_delete += 1
    with tqdm(total=num_to_delete) as progress:
        for table in all_tables:
            if "integration_test" in table.name:
                table.delete()
                progress.update()
    print(f"Deleted {num_to_delete} CI DynamoDB tables")


def cleanup_bigtable_ci():
    client = bigtable.Client(project="kf-feast", admin=True)
    instance = client.instance("feast-integration-tests")
    if instance.exists():
        print(f"Deleted Bigtable CI instance")
        instance.delete()

    location_id = "us-central1-f"
    serve_nodes = 1
    storage_type = enums.StorageType.SSD
    cluster = instance.cluster(
        "feast-integration-tests-c1",
        location_id=location_id,
        serve_nodes=serve_nodes,
        default_storage_type=storage_type,
    )
    instance.create(clusters=[cluster])
    print(f"Created new Bigtable CI tables")


def main() -> None:
    cleanup_dynamo_ci()
    cleanup_bigtable_ci()


if __name__ == "__main__":
    main()