#
# n23 - data acquisition and processing framework
#
# Copyright (C) 2013-2025 by Artur Wroblewski <wrobell@riseup.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

"""
Implementation of default data storage for N23 framework.

The storage is backed by PostgreSQL RDBMS. It is tested with TimescaleDB
and PostGIS extensions.

The storage is implemented with :py:class:`Database` class.

A database table, which stores application data, needs to be declared with
:py:meth:`Database.add_entity` method. It also allows to specify table
columns.

For example::

    db.add_entity('temperature`, ('location', 'value'))

and for task results::

    TaskResult(time=10001, name='temperature/room1', value=20.1)
    TaskResult(time=10002, name='temperature/room1', value=20.2)
    TaskResult(time=10003, name='temperature/room1', value=20.3)

the data is stored in table `temperature` as::

    +-------+----------+-------+
    | time  | location | value |
    +-------+----------+-------+
    | 10001 | room1    |  20.1 |
    | 10002 | room1    |  20.2 |
    | 10003 | room1    |  20.3 |
    +-------+----------+-------+

"""

import abc
import asyncio
import asyncpg  # type: ignore
import logging
import types
import typing as tp
from collections import deque
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from functools import singledispatch, partial
from urllib.parse import urlparse, urlunparse

from ..api import DEFAULT_MAX_LEN, Storage, Fields
from ...fn import identity
from ...scheduler import TaskResult
from ...util import to_datetime

logger = logging.getLogger(__name__)

T_co = tp.TypeVar('T_co', covariant=True)
TypeEncoder: tp.TypeAlias = tp.Callable[..., tp.Any]
TypeEncoders: tp.TypeAlias = dict[str, TypeEncoder]
TableRow: tp.TypeAlias = tuple[tp.Any, ...]

# see https://github.com/MagicStack/asyncpg/issues/513
CONNECTION_ERRORS = (
    OSError,
    asyncpg.CannotConnectNowError,
    asyncpg.ConnectionDoesNotExistError,
)

CONNECTION_ERROR_BACK_TO_POOL = 'connection has been released back to the pool'
CONNECTION_ERROR_CLOSED = 'the underlying connection is closed'

SLEEP_ON_CONNECTION_ERROR = 5

class Database(Storage):
    """
    N23 framework storage implementation using PostgreSQL RDBMS.
    """
    def __init__(self, db_uri: str) -> None:
        """
        Initialize storage backend.

        :param db_uri: Database URI specification.
        """
        self.db_uri = db_uri

        # connection pool size controls how many parallel connections are
        # performed
        self.pool_size = 1

        self._type_codecs: TypeEncoders = {}
        self._writers: dict[str, TableWriterBase] = {}

        self._pool = asyncpg.create_pool(
            self.db_uri,
            min_size=1,
            max_size=self.pool_size,
            # TODO: use context vars instead of partial application
            init=partial(set_connection_codecs, self._type_codecs)
        )

    def add_entity(
            self,
            name: str,
            fields: Fields,
            *,
            batch: bool=False,
            max_len: int=DEFAULT_MAX_LEN,
    ) -> None:
        """
        Add and configure storage entity - PostgreSQL database table.
        """
        if name in self._writers:
            raise ValueError('Entity "{}" already defined' .format(name))
        if len(fields) < 1:
            raise ValueError('At least one column expected')

        cols = ('time', *fields)
        writer: TableWriterBase
        if batch:
            writer = TableBatchWriter(self._pool, name, cols, max_len)
        else:
            writer = TableWriter(self._pool, name, cols)

        self._writers[name] = writer

    async def write(self, result: TaskResult[T_co]) -> None:
        """
        Write N23 framework task result into database table.

        :param result: Task result generated by N23 scheduler.
        """
        table, *name_values = result.name.split('/')
        ts = to_datetime(result.time)
        writer = self._writers[table]

        # is single or multiple values expected in task result?
        # compare number of table declared columns with number of values to
        # be inserted into the table
        # - time column
        # - values extracted from the task result name
        # - at least one value from the task result
        is_single = len(writer._columns) == len(name_values) + 2  # noqa: SLF001
        values = (result.value,) if is_single else result.value

        row = (ts, *name_values, *values)  # type: ignore
        await writer.write(row)

    def set_type_codec(self, type_name: str, encoder: TypeEncoder) -> None:
        """
        Assign encoder to a database type.

        :param type_name: Database type name.
        :param encoder: Function converting Python value to format understood
            by database.
        """
        self._type_codecs[type_name] = encoder

    async def flush(self) -> None:
        tasks = (
            w.flush() for w in self._writers.values()
            if hasattr(w, 'flush')
        )
        task = asyncio.gather(*tasks)
        await catch_conn_error(task)

    async def __aenter__(self) -> None:
        await self._pool

    async def __aexit__(
        self,
        ex_type: type[BaseException] | None,
        ex: BaseException | None,
        traceback: types.TracebackType | None,
    ) -> None:
        try:
            async with asyncio.timeout(30):
                await self.flush()
                await self._pool.close()
        except asyncio.TimeoutError:
            logger.warning('database flush failed due to timeout: %s', ex)

class TableWriterBase:
    """
    Basic class for PostgreSQL database table writers.
    """
    def __init__(
            self,
            pool: asyncpg.Pool,
            table: str,
            columns: Fields,
    ):
        self._pool = pool
        self._table = table
        self._columns = columns

    @abc.abstractmethod
    async def write(self, row: TableRow) -> None:
        ...

class TableWriter(TableWriterBase):
    """
    PostgreSQL database table writer.
    """
    async def write(self, row: TableRow) -> None:
        """
        Writer row to a PostgreSQL database table.

        :param row: Data to write.
        """
        while True:
            task = self._write(row)
            saved = await catch_conn_error(task)
            if saved:
                break
            else:
                await asyncio.sleep(SLEEP_ON_CONNECTION_ERROR)

    async def _write(self, row: TableRow) -> None:
        async with connect(self._pool) as conn:  # type: ignore
            await conn.copy_records_to_table(
                self._table,
                records=[row],
                columns=self._columns,
            )

        if __debug__:
            logger.debug('stored 1 record in table {}'.format(self._table))

class TableBatchWriter(TableWriterBase):
    """
    PostgreSQL database table batch writer.

    Data is added to an internal queue on write. Data is flushed to
    database table using :py:meth:`TableBatchWriter.flush` method.
    """
    def __init__(
            self,
            pool: asyncpg.Pool,
            table: str,
            columns: Fields,
            max_len: int,
    ):
        super().__init__(pool, table, columns)
        self._queue = deque[TableRow]([])
        self._lock = asyncio.Lock()
        self._max_len = max_len

    async def write(self, row: TableRow) -> None:
        """
        Enqueue row of data to be written into a PostgreSQL database table.

        :param row: Data to write.
        """
        # use lock to avoid error, when asyncpg iterates over the table
        # rows while saving into database
        async with self._lock:
            self._queue.append(row)

        if len(self._queue) == self._max_len:
            await self.flush()

    async def flush(self) -> None:
        """
        Flush enqueued rows of data into a PostgreSQL database table.

        Do nothing if internal queue is empty.
        """
        if not self._queue:
            return

        async with self._lock, connect(self._pool) as conn:  # type: ignore
            await conn.copy_records_to_table(
                self._table,
                records=self._queue,
                columns=self._columns,
            )
            count_rows = len(self._queue)
            self._queue.clear()

        if __debug__:
            logger.debug(
                'stored {} record(s) in table {}'.format(count_rows, self._table)
            )

async def catch_conn_error(coro: tp.Awaitable[tp.Any]) -> bool:
    """
    Protect calling coroutine from connection error of coroutine being
    executed.

    Return true if coroutine executed without connection error.

    :param coro: Coroutine to execute.
    """
    try:
        await coro
    except CONNECTION_ERRORS as ex:
        logger.warning(
            'connection failure due to: {}'.format(ex)
        )
        return False
    except asyncpg.InterfaceError as ex:
        if not is_disconnected(ex):
            raise
        logger.warning(
            'connection failure due to: {}'.format(ex)
        )
        return False
    else:
        return True

@asynccontextmanager
async def create_pool(
        db_uri: str,
        type_codecs: TypeEncoders={},
        max_size: int=1,
    ) -> AsyncIterator[asyncpg.Pool]:
    """
    Create connection pool for a database.

    :param db_uri: Database URI.
    :param type_codecs: Mapping between database types and their encoders.
    :param max_size: Maximum number of connections of connection pool.
    """
    async with asyncpg.create_pool(
                db_uri,
                min_size=1,
                max_size=max_size,
                init=partial(set_connection_codecs, type_codecs),
            ) as pool:

        yield pool

@singledispatch  # type: ignore
@asynccontextmanager
async def connect(
        db: str | asyncpg.Pool,
        type_codecs: TypeEncoders={},
        max_size: int=1,
        enable_tx: bool=True
    ) -> AsyncIterator[asyncpg.Connection]:
    """
    Context manager to connect to a database.

    First parameter can be

    - database URI
    - or database connection pool

    In first instance, new connection pool is created to connect to
    a database. In second instance, the connection pool is reused to
    connect to a database.

    :param db: Database URI or database connection pool.
    :param max_size: Maximum number of connections of connection pool.
    :param enable_tx: Create database transaction (true by default).
    """
    raise NotImplementedError(
        'Not implemented for database parameter type: {}'
        .format(type(db))
    )

@connect.register
@asynccontextmanager
async def _connect_db_uri(
        db_uri: str,
        type_codecs: TypeEncoders={},
        max_size: int=1,
        enable_tx: bool=True
    ) -> AsyncIterator[asyncpg.Connection]:
    """
    Create connection pool and use it to connect to a database.
    """
    async with create_pool(db_uri, type_codecs, max_size=max_size) as pool:
        async with connect(pool, enable_tx=enable_tx) as conn:  # type: ignore
            yield conn

@connect.register
@asynccontextmanager
async def _connect_db_pool(
        pool: asyncpg.Pool,
        enable_tx: bool=True
    ) -> AsyncIterator[asyncpg.Connection]:
    """
    Reuse existing connection pool to connect to a database.
    """
    async with pool.acquire() as conn:
        if enable_tx:
            async with conn.transaction():
                yield conn
        else:
            yield conn

async def set_connection_codecs(
        type_codecs: TypeEncoders,
        conn: asyncpg.Connection,
    ) -> None:
    set_codec = partial(
        conn.set_type_codec, decoder=identity, format='binary'
    )
    for t, f in type_codecs.items():
        if __debug__:
            logger.debug('set codec for type {}: {}'.format(t, f))
        await set_codec(t, encoder=f)

async def execute(db_uri: str, cmd: str, enable_tx: bool=True) -> None:
    """
    Connect to database and execute single database statement.
    """
    logger.info('execute: {}'.format(cmd))
    async with connect(db_uri, enable_tx=enable_tx) as conn:  # type: ignore
        await conn.execute(cmd)

def is_disconnected(ex: asyncpg.InterfaceError) -> bool:
    """
    Check if the PostgreSQL interface error is a disconnection error.
    """
    message = str(ex)
    disconnected = CONNECTION_ERROR_BACK_TO_POOL in message \
        or CONNECTION_ERROR_CLOSED in message
    logger.debug(
        'interface error (disconnected={}): {}'
        .format(disconnected, message)
    )
    return disconnected

def to_db_uri(db: str, alt_db: tp.Optional[str]=None) -> str:
    """
    Convert database name or partial database connection string to database
    connection string for local database server.

    If the `db` parameter is database connection string already, then do
    nothing.

    For example::

        some-db -> postgresql://localhost/some-db
        localhost/some-db -> postgresql://localhost/some-db
        postgresql://localhost/some-db -> postgresql://localhost/some-db
        user:pass@localhost/some-db -> postgresql://user:pass@localhost/some-db

    Use alternative database name to get connection for template database,
    i.e. `template1`.

    @param db: Database name or database connection string.
    @param alt_db: Alternative database name.
    """
    db = db if db.startswith('postgresql://') or db.startswith('//') else '//' + db
    db_uri = urlparse(db, scheme='postgresql')

    # if no hostname in `db`, then netloc points to dbname, therefore make
    # a swap and use localhost for hostname
    host, dbname = (db_uri.netloc, db_uri.path) \
        if db_uri.netloc and db_uri.path \
        else ('localhost', db_uri.netloc)
    dbname = dbname if alt_db is None else alt_db

    db_uri = db_uri._replace(path=dbname, netloc=host)
    logger.info('parsed db uri: {}'.format(db_uri))
    return urlunparse(db_uri)

# vim: sw=4:et:ai
