import asyncio
import time
from contextlib import ExitStack, AsyncExitStack
from functools import partialmethod

import pytest

from clarabot_nano.utility.aio_sqlite import AbstractAioSQLiteDatabase
from clarabot_nano_tests.testing.file import temp_path


class _SQLiteDbgException(Exception):
    pass


class _TestingAioSQLiteDatabase(AbstractAioSQLiteDatabase):
    def _initialise(self):
        self._con.execute("PRAGMA journal_mode = WAL;")
        self._con.execute("PRAGMA synchronous = NORMAL;")

    def _setup(self):
        self._con.execute("CREATE TABLE IF NOT EXISTS Data0 (key TEXT PRIMARY KEY, value TEXT)")

    def _write_impl(self, key, value, error=None, delay=0.0):
        time.sleep(delay)
        if error:
            raise error
        c = self._con.cursor()
        c.execute("REPLACE INTO Data0(key, value) VALUES (?, ?)", (key, value))
        return c.lastrowid

    def _read_impl(self, key):
        rows = self._con.execute("SELECT value FROM Data0 WHERE key = ?", (key,)).fetchall()
        if rows:
            return rows[0][0]

    write = partialmethod(AbstractAioSQLiteDatabase._run_job, _write_impl)
    write_no_immediate = partialmethod(AbstractAioSQLiteDatabase._run_job, _write_impl, begin_immediate=False)
    read = partialmethod(AbstractAioSQLiteDatabase._run_job, _read_impl, begin_immediate=False)


@pytest.mark.asyncio
async def test_aio_sqlite_wal_retry(tmpdir):
    loop = asyncio.get_event_loop()
    db_path = str(temp_path(tmpdir))
    async with AsyncExitStack() as stack:
        dba = await stack.enter_async_context(_TestingAioSQLiteDatabase.create(loop, db_path, 10.0))
        dbb = await stack.enter_async_context(_TestingAioSQLiteDatabase.create(loop, db_path, 10.0))
        dbc = await stack.enter_async_context(_TestingAioSQLiteDatabase.create(loop, db_path, 10.0))
        await asyncio.sleep(0.25)  # all need to start up to have stable locking

        await dba.write("kx", "vx")
        assert await dbc.read("kx") == "vx"

        t = time.perf_counter()

        async def _read_before(db_, k_, r_):
            v_ = await db_.read(k_)
            assert time.perf_counter() < t
            assert v_ == r_

        # test multiple writes serialize and allow parallel reads
        t = time.perf_counter() + 0.5
        at = asyncio.create_task(dba.write("k0", "v0", None, 0.5))
        await asyncio.sleep(0.1)
        await _read_before(dbc, "kx", "vx")
        await _read_before(dbc, "k0", None)
        bt = asyncio.create_task(dbb.write("k1", "v1", None, 0.025))
        await asyncio.sleep(0.1)
        await _read_before(dbc, "kx", "vx")
        await _read_before(dbc, "k0", None)
        await _read_before(dbc, "k1", None)
        await at
        assert not bt.done()
        assert time.perf_counter() >= t
        await bt
        assert await dbc.read("kx") == "vx"
        assert await dbc.read("k0") == "v0"
        assert await dbc.read("k1") == "v1"

        # tx fail is handled well
        t = time.perf_counter() + 0.5
        at = asyncio.create_task(dba.write("k0", "v2", _SQLiteDbgException(), 0.5))
        await asyncio.sleep(0.1)
        bt = asyncio.create_task(dbb.write("k0", "v3", None, 0.5))
        await asyncio.sleep(0.1)
        await _read_before(dbc, "k0", "v0")
        with pytest.raises(_SQLiteDbgException):
            await at
        t += 0.5
        await _read_before(dbc, "k0", "v0")
        await bt
        assert time.perf_counter() >= t
        assert await dbc.read("kx") == "vx"
        assert await dbc.read("k0") == "v3"
        assert await dbc.read("k1") == "v1"

        # test non-immediate write behaviour
        t = time.perf_counter() + 0.5
        at = asyncio.create_task(dba.write("k0", "v4", None, 0.5))
        await asyncio.sleep(0.1)
        await _read_before(dbc, "k0", "v3")
        await dbb.write_no_immediate("k0", "v5")  # deferred waits without retry loop
        assert at.done()
        assert time.perf_counter() >= t
        assert await dbc.read("kx") == "vx"
        assert await dbc.read("k0") == "v5"
        assert await dbc.read("k1") == "v1"
