# retoor <retoor@molodetz.nl>

from __future__ import annotations

import asyncio
from collections.abc import Callable
from pathlib import Path
from typing import Any, TypeVar

import dataset

from devplacepy_services.base.config import sqlite_read_pool_size

T = TypeVar("T")

PRAGMAS = [
    "PRAGMA journal_mode=WAL",
    "PRAGMA synchronous=NORMAL",
    "PRAGMA busy_timeout=30000",
    "PRAGMA cache_size=-8000",
    "PRAGMA temp_store=MEMORY",
    "PRAGMA mmap_size=268435456",
]


class SQLiteFileBroker:
    name: str
    path: Path
    read_pool_size: int

    def __init__(self, name: str, path: Path, read_pool_size: int | None = None) -> None:
        self.name = name
        self.path = Path(path)
        self.read_pool_size = read_pool_size or sqlite_read_pool_size()
        self._write_db: Any = None
        self._read_pool: asyncio.Queue[Any] | None = None
        self._write_queue: asyncio.Queue[Any] | None = None
        self._write_worker_task: asyncio.Task | None = None

    def _connect_rw(self):
        return dataset.connect(
            f"sqlite:///{self.path}",
            engine_kwargs={
                "connect_args": {
                    "timeout": 30,
                    "check_same_thread": False,
                },
            },
            on_connect_statements=PRAGMAS,
        )

    def _connect_ro(self):
        abs_path = self.path.resolve().as_posix()
        return dataset.connect(
            f"sqlite:///file:{abs_path}?uri=true&mode=ro",
            engine_kwargs={
                "connect_args": {
                    "timeout": 30,
                    "check_same_thread": False,
                    "uri": True,
                },
            },
            on_connect_statements=PRAGMAS,
        )

    async def startup(self) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._write_db = self._connect_rw()
        self._read_pool = asyncio.Queue()
        for _ in range(self.read_pool_size):
            await self._read_pool.put(self._connect_ro())
        self._write_queue = asyncio.Queue()
        self._write_worker_task = asyncio.create_task(self._write_worker())

    async def shutdown(self) -> None:
        if self._write_worker_task is not None:
            self._write_worker_task.cancel()
            try:
                await self._write_worker_task
            except asyncio.CancelledError:
                pass
            self._write_worker_task = None

    async def read(self, fn: Callable[..., T], *args, **kwargs) -> T:
        assert self._read_pool is not None
        db = await self._read_pool.get()
        try:
            return await asyncio.to_thread(fn, db, *args, **kwargs)
        finally:
            await self._read_pool.put(db)

    async def write(self, fn: Callable[..., T], *args, **kwargs) -> T:
        assert self._write_queue is not None
        loop = asyncio.get_running_loop()
        future = loop.create_future()
        await self._write_queue.put((fn, args, kwargs, future))
        return await future

    async def _write_worker(self) -> None:
        assert self._write_queue is not None
        while True:
            fn, args, kwargs, future = await self._write_queue.get()
            try:
                result = fn(self._write_db, *args, **kwargs)
                future.set_result(result)
            except Exception as exc:
                future.set_exception(exc)

    async def read_after_write(self, fn: Callable[..., T], *args, **kwargs) -> T:
        return await asyncio.to_thread(fn, self._write_db, *args, **kwargs)