import asyncio
from multiprocessing import Process, Queue
from multiprocessing.queues import Empty
from typing import Any, Callable, List

import aiomysql


class AsyncWorker:
    # to avoid concurrent connect
    _connection_lock = asyncio.Lock()

    def __init__(self, queue, async_connection_provider):
        self._connections = {}
        self._queue = queue
        self.async_connection_provider = async_connection_provider

    async def get_connection(self, schema_name: str) -> aiomysql.Connection:
        if schema_name not in self._connections:
            async with self._connection_lock:
                self._connections[schema_name] = await self.async_connection_provider(
                    db=schema_name
                )
        return self._connections[schema_name]

    async def reset_connection(self, schema_name: str) -> aiomysql.Connection:
        self._connections[schema_name].close()
        self._connections[schema_name] = await self.async_connection_provider(
            db=schema_name
        )
        return self._connections[schema_name]

    async def run(self):
        while True:
            try:
                msg = self._queue.get(block=False)
            except Empty:
                await asyncio.sleep(0.1)
                continue

            if msg is None:
                break
            schema, statement, parameters = msg
            conn = await self.get_connection(schema)

            try:
                async with conn.cursor() as cur:
                    await cur.execute(statement, parameters)
                    await conn.commit()
            except Exception:
                # TODO: add exception into report
                await self.reset_connection(schema)


def worker(
    queue: Queue,
    concurrency,
    async_connection_provider,
):
    async def _worker_impl():
        ws = [AsyncWorker(queue, async_connection_provider) for _ in range(concurrency)]
        await asyncio.gather(*[asyncio.create_task(w.run()) for w in ws])

    asyncio.run(_worker_impl())


class AsyncExecutor:
    def __init__(
        self,
        concurrency: int,
        async_connection_provider: Callable[..., aiomysql.Connection],
        async_queue_size: int = 1000,
    ):
        self.concurrency = concurrency
        print(async_queue_size)
        self.queue = Queue(async_queue_size)
        self.process = Process(
            target=worker, args=(self.queue, concurrency, async_connection_provider)
        )
        self.process.start()

    def execute(self, schema: str, statement: str, parameters: List[Any]):
        self.queue.put((schema, statement, parameters))

    def run(self, concurrency: int):
        asyncio.run(self.async_replay(concurrency))

    def kill(self):
        for _ in range(self.concurrency):
            self.queue.put(None)
        self.process.join()