"""mysql connector."""

from contextlib import contextmanager
from typing import Any, Generator

from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine
from sqlalchemy.orm import DeclarativeMeta, Session, declarative_base, sessionmaker

from carveouts import config


class SqlConnector:
    """SQL connector."""

    base_model: DeclarativeMeta
    _db_engine: Engine
    _db_session_maker: sessionmaker[Any]

    def __init__(self, db_url: str) -> None:
        self.base_model = declarative_base()
        self._db_engine = create_engine(
            db_url,
            **config.DB_CONNECT_ARGS,
        )
        self._db_session_maker = sessionmaker(
            bind=self._db_engine,
            class_=Session,
            expire_on_commit=False,
        )

    @contextmanager
    def db_session(
        self, transaction: bool = False, turn_off_foreign_key_constraint: bool = False
    ) -> Generator[Session, None, None]:
        query_error = None
        cleanup_error = None
        session = self._db_session_maker()

        # do not use transaction
        if not transaction:
            session.connection(
                execution_options={"isolation_level": "AUTOCOMMIT", "autobegin": False}
            )

        # give back session to query with
        try:
            if turn_off_foreign_key_constraint:
                session.execute(text("SET FOREIGN_KEY_CHECKS = 0;"))
            yield session
            if transaction:
                session.commit()
        except Exception as e:
            query_error = e
            if transaction:
                session.rollback()
        finally:
            if turn_off_foreign_key_constraint:
                session.execute(text("SET FOREIGN_KEY_CHECKS = 1;"))

        # cleanup by closing session
        try:
            session.close()
        except Exception as e:
            cleanup_error = e

        # report errors
        if query_error:
            raise query_error
        if cleanup_error:
            raise cleanup_error

    def close(self) -> None:
        self._db_engine.dispose()

    def engine(self) -> Engine:
        return self._db_engine


db_connector = SqlConnector(config.AR_DB_URL)