from minio import Minio
import os

from constants import MINIO_SERVER, MINIO_DATA_LABELED, MINIO_MODEL_TRAINED
from logger import logger

minio_client = Minio(
    endpoint=MINIO_SERVER,
    access_key=os.getenv("MINIO_ROOT_USER"),
    secret_key=os.getenv("MINIO_ROOT_PASSWORD"),
    secure=False
)
bucket_names = [MINIO_DATA_LABELED, MINIO_MODEL_TRAINED]

def check_bucket(minio_client=minio_client, bucket_names=bucket_names):
    # Check if bucket exists
    for bucket_name in bucket_names:
        found = minio_client.bucket_exists(bucket_name)
        if not found:
            # Create bucket
            minio_client.make_bucket(bucket_name)
            logger.info(f"Bucket '{bucket_name}' created successfully.")
        else:
            logger.info(f"Bucket '{bucket_name}' already exists.")