import logging
import uuid
import time
from sqlalchemy.exc import SQLAlchemyError, OperationalError

from app.settings import settings
from app.utils.tidb_serverless import TiDBServerlessClient
from app.db import Scoped_Session
from app.repositories.db_instance import DbInstanceRepo

logging.basicConfig(level=settings.LOG_LEVEL)
logger = logging.getLogger(__name__)

tcs = TiDBServerlessClient(
    public_key=settings.TiDBCLOUD_PUBLIC_KEY,
    private_key=settings.TiDBCLOUD_PRIVATE_KEY,
    project_id=settings.TiDBCLOUD_PROJECT_ID,
    cluster_id=settings.TiDBCLOUD_CLUSTER_ID
)

db_repo = DbInstanceRepo(tcs)

pool_size = 20

if __name__ == '__main__':
    while True:
        with Scoped_Session() as session:
            try:
                db_repo.clear_branch_instances(session)
                instances = db_repo.get_all_available_db_instance(session)
                logger.info(f"Available instances: {len(instances)}")
                if len(instances) <= pool_size:
                    branch_name = str(uuid.uuid4())
                    db_repo.create_branch_instance(session, branch_name)
                    logger.info(f"Branch instance created: {branch_name}")
            except OperationalError as oe:
                session.rollback()
                logger.error(f"OperationalError occurred: {oe}")
            except SQLAlchemyError as e:
                session.rollback()
                logger.error(f"SQLAlchemyError occurred: {e}")
            except Exception as e:
                logger.error(f"Error occurred: {e}")
                session.rollback()

        time.sleep(20)