from msgbus import MBEasyContext
from persistence import (
    Driver,
    Connection,
    Cursor,
)

from typing import Optional


class ConnectionWrapper(object):
    def __init__(
        self,
        mb_context,         # type: MBEasyContext
        db_name=None,       # type: Optional[str]
        service_name=None,  # type: Optional[str]
    ):
        self.mb_context = mb_context
        self.db_name = db_name
        self.service_name = service_name
        self.conn = None  # type: Optional[Connection]

    def query(
        self,   # type: ConnectionWrapper
        sql,    # type: str
    ):
        # type: (...) -> int

        return self._get_conn().query(sql=sql)

    def select(
        self,   # type: ConnectionWrapper
        sql,    # type: str
    ):
        # type: (...) -> Cursor

        return self._get_conn().select(sql=sql)

    def begin_transaction(
        self,   # type: ConnectionWrapper
    ):
        # type: (...) -> None

        conn = self._get_conn()
        conn.transaction_start()
        conn.query(sql="BEGIN TRANSACTION")

    def commit(
        self,   # type: ConnectionWrapper
    ):
        # type: (...) -> None

        self.conn.query(sql="COMMIT")
        self.conn.transaction_end()

    def rollback(
        self,   # type: ConnectionWrapper
    ):
        # type: (...) -> None

        self.conn.query(sql="ROLLBACK")
        self.conn.transaction_end()

    def close(
        self,   # type: ConnectionWrapper
    ):
        # type: (...) -> None

        if self.conn is not None:
            self.conn.close()

    def _get_conn(
        self,   # type: ConnectionWrapper
    ):
        # type: (...) -> Connection

        if self.conn is None:
            self.conn = Driver().open(
                mbcontext=self.mb_context,
                dbname=str(self.db_name.encode("utf-8")) if self.db_name is not None else None,
                service_name=self.service_name.encode("utf-8"),
            )
        return self.conn
