import asyncio
import sqlite3
from contextlib import contextmanager
from functools import partial
import queue
from typing import Callable, Optional

import wait_for2

from clarabot_nano.utility.thread import AbstractThread


def _callback_result(future: asyncio.Future, result):
    if not future.done():
        future.set_result(result)


def _callback_exception(future: asyncio.Future, exc):
    if not future.done():
        future.set_exception(exc)


class AbstractAioSQLiteDatabase(AbstractThread):
    """
    Provide a lightweight asyncio compatible, customizable interface to an arbitrary SQLite database in a safe way.
    All SQL connection operations are restricted to be executed on the worker thread, guaranteeing serialization
    requirements. If more processes would access the database the transaction_mode selector can be used, but job
    specific transaction handling would be needed for optimal performance.

    Requests shall be implemented like this:

        def _get_all_impl(self):
            return self._con.execute("SELECT x FROM y").fetchall()

        get_all = partialmethod(AbstractAioSQLiteDatabase._run_job, _get_all_impl)

    """

    @classmethod
    @contextmanager
    def create(cls, *args, **kwargs):
        """
        Strict start-stop helper for automatic cleanup.
        """
        self = cls(*args, **kwargs)
        try:
            self.start()
            yield self
        finally:
            self.stop()

    def __init__(
        self,
        loop: asyncio.AbstractEventLoop,
        sqlite_db_path: str,
        sqlite_timeout: float = 60.0,
        transaction_mode: Optional[str] = None,
    ):
        self._loop = loop
        self._con: sqlite3.Connection = None
        self._request_queue: queue.Queue = None
        self._begin_statement = self.parse_begin_statement(transaction_mode)
        AbstractThread.__init__(self, sqlite_db_path, sqlite_timeout)
        self.closed = asyncio.Event()

    @staticmethod
    def parse_begin_statement(transaction_mode: Optional[str]) -> str:
        if transaction_mode is None or transaction_mode == "DEFERRED":
            return "BEGIN"
        elif transaction_mode in ("IMMEDIATE", "EXCLUSIVE"):
            return f"BEGIN {transaction_mode}"
        raise ValueError(f"Invalid transaction mode {transaction_mode}")

    def start(self, callback=None, daemon=False, name_prefix=None):
        if callback is not None:  # pragma: no cover
            raise ValueError("SQLite worker does not support custom callback. Use the 'closed' event.")
        if daemon not in (None, False):  # pragma: no cover
            raise ValueError("SQLite worker must not be daemon to ensure cleanup.")
        self._request_queue = queue.Queue()
        AbstractThread.start(self, partial(self._loop.call_soon_threadsafe, self.closed.set), False, name_prefix)

    def stop(self, timeout=5.0):
        q = self._request_queue
        if q is not None:
            q.put_nowait(None)
        AbstractThread.stop(self, timeout)

    def _thread(self, db_path: str, timeout: float):
        """
        SQLite connection and request executor thread.
        """
        self._con = sqlite3.connect(db_path, timeout=timeout, isolation_level=None)
        try:
            # initialise the db
            self._initialise()
            with self._transact("BEGIN IMMEDIATE" if self._begin_statement == "BEGIN" else self._begin_statement):
                self._setup()
            # run job queue
            while True:
                job = self._request_queue.get()
                if job is None:
                    break
                process_fn, args, post_process_fn, aio_future = job  # TODO: add tr-begin preference to job for optimization?
                if aio_future.done():
                    continue
                try:
                    with self._transact(self._begin_statement):
                        result = process_fn(self, *args)
                    if post_process_fn:
                        post_process_fn(self)
                    self._loop.call_soon_threadsafe(_callback_result, aio_future, result)
                except Exception as e:
                    self._loop.call_soon_threadsafe(_callback_exception, aio_future, e)
        finally:
            # prevent more jobs to be queued
            q = self._request_queue
            self._request_queue = None
            # notify dead jobs if any
            try:
                while True:
                    job = q.get_nowait()
                    if job:
                        job[3].cancel()
            except queue.Empty:
                pass
            # close the database
            con = self._con
            self._con = None
            con.close()

    @contextmanager
    def _transact(self, begin_statement: str):
        self._con.execute(begin_statement)
        try:
            yield
        except BaseException:
            self._con.execute("ROLLBACK")
            raise
        else:
            self._con.execute("COMMIT")

    def _create_job(self, process_fn, args, post_process_fn) -> asyncio.Future:
        """
        This will raise AttributeError if the database has been closed.
        """
        self._request_queue.put_nowait((process_fn, args, post_process_fn, f := self._loop.create_future()))
        return f

    async def _run_job(
        self, process_fn: Callable, *args, post_process_fn: Optional[Callable] = None, timeout: Optional[float] = None
    ):
        """
        This is designed to be used as

            get_xy = partialmethod(AbstractAioSQLiteDatabase._run_job, _get_xy_impl)

        which means the wrapped function is not yet bound in the expression. The self argument is manually passed by
        this method to make this work.
        """
        return await wait_for2.wait_for(self._create_job(process_fn, args, post_process_fn), timeout)

    def _initialise(self):
        """
        The connection is ready and the database initialization like pragma definitions shall be done here.
        NOTE: Current execution is not inside a transaction!
        """

    def _setup(self):
        """
        The connection is ready and the database initialization like schema shall be done.
        """
        raise NotImplementedError()
