import typing

import databases
from aioworkers.core.base import AbstractConnector
from aioworkers.core.config import Config
from databases.core import Connection, Transaction
from sqlalchemy.sql import ClauseElement


class Database(AbstractConnector):
    """
    Use Wrapper (Decorator) pattern to wrap databases.Database and add it behavior to entity.
    """

    def __init__(self, config=None, *, context=None, loop=None):
        super().__init__(config, context=context, loop=loop)
        self._db: databases.Database = None
        self._dsn: str = None

    async def init(self):
        await super().init()
        self._db = databases.Database(self._dsn)

    def set_config(self, config: Config) -> None:
        cfg = config.new_parent(logger="aioworkers_databases")
        self._dsn = cfg.get("dsn")
        super().set_config(cfg)

    async def connect(self):
        await self._db.connect()

    async def disconnect(self):
        await self._db.disconnect()

    async def execute(
        self, query: typing.Union[ClauseElement, str], values: dict = None
    ) -> typing.Any:
        return await self._db.execute(query, values)

    async def execute_many(
        self, query: typing.Union[ClauseElement, str], values: list
    ) -> None:
        return await self._db.execute_many(query, values)

    async def fetch_all(
        self, query: typing.Union[ClauseElement, str], values: dict = None
    ) -> typing.List[typing.Mapping]:
        return await self._db.fetch_all(query, values)

    async def fetch_one(
        self, query: typing.Union[ClauseElement, str], values: dict = None
    ) -> typing.Optional[typing.Mapping]:
        return await self._db.fetch_one(query, values)

    async def fetch_val(
        self,
        query: typing.Union[ClauseElement, str],
        values: dict = None,
        column: typing.Any = 0,
    ) -> typing.Any:
        return await self._db.fetch_val(query, values, column)

    def transaction(
        self, *, force_rollback: bool = False, **kwargs: typing.Any
    ) -> "Transaction":
        return self._db.transaction(force_rollback=force_rollback, **kwargs)

    def iterate(
        self, query: typing.Union[ClauseElement, str], values: dict = None
    ) -> typing.AsyncGenerator[typing.Mapping, None]:
        return self._db.iterate(query, values)

    def connection(self) -> "Connection":
        return self._db.connection()
