ytptube/app/library/DataStore.py

import asyncio
import copy
import logging
from collections import OrderedDict
from collections.abc import Iterable
from enum import Enum

from .Download import Download
from .ItemDTO import ItemDTO
from .operations import matches_condition
from .sqlite_store import SqliteStore

LOG = logging.getLogger("datastore")


class StoreType(str, Enum):
    HISTORY = "done"
    QUEUE = "queue"

    @classmethod
    def all(cls) -> list[str]:
        return [member.value for member in cls]

    @classmethod
    def from_value(cls, value: str) -> "StoreType":
        for member in cls:
            if member.value == value:
                return member
        msg = f"Invalid StoreType value: {value}"
        raise ValueError(msg)

    def __str__(self) -> str:  # pragma: no cover
        return self.value


def _strip_transient_fields(item: ItemDTO) -> ItemDTO:
    stored = copy.deepcopy(item)
    if hasattr(stored, "datetime"):
        try:
            delattr(stored, "datetime")
        except AttributeError:
            pass
    if stored.status == "finished" and hasattr(stored, "live_in"):
        try:
            delattr(stored, "live_in")
        except AttributeError:
            pass
    return stored


class DataStore:
    """
    In-memory queue cache + optional write-behind to persistence.
    History reads go straight to persistence.
    """

    def __init__(self, type: StoreType, connection: SqliteStore):
        self._type = type
        self._connection: SqliteStore = connection
        self._dict: OrderedDict[str, Download] = OrderedDict()

    # cache helpers
    async def load(self) -> None:
        saved = await self._connection.fetch_saved(str(self._type))
        for key, item in saved:
            self._dict[key] = Download(info=item)

    async def saved_items(self) -> list[tuple[str, ItemDTO]]:
        return await self._connection.fetch_saved(str(self._type))

    def exists(self, key: str | None = None, url: str | None = None) -> bool:
        if not key and not url:
            msg = "key or url must be provided."
            raise KeyError(msg)
        if key and key in self._dict:
            return True
        return any(
            (key and self._dict[i].info._id == key) or (url and self._dict[i].info.url == url) for i in self._dict
        )

    def get(self, key: str | None = None, url: str | None = None) -> Download:
        if not key and not url:
            msg = "key or url must be provided."
            raise KeyError(msg)
        for i in self._dict:
            if (key and self._dict[i].info._id == key) or (url and self._dict[i].info.url == url):
                return self._dict[i]
        msg: str = f"{key=} or {url=} not found."
        raise KeyError(msg)

    def get_item(self, **kwargs) -> Download | None:
        if not kwargs:
            return None
        for i in self._dict:
            if not self._dict[i].info:
                continue
            info = self._dict[i].info.__dict__
            if any(matches_condition(key, value, info) for key, value in kwargs.items()):
                return self._dict[i]
        return None

    async def get_by_id(self, id: str) -> Download | None:
        val = self._dict.get(id)
        if val or self._type == StoreType.QUEUE:
            return val

        if item := await self._connection.get_by_id(str(self._type), id):
            download = Download(info=item)
            self._dict[id] = download
            return download

        return None

    def items(self):
        return self._dict.items()

    async def put(self, value: Download, no_notify: bool = False) -> Download:
        _ = no_notify
        self._dict[value.info._id] = value
        asyncio.create_task(self._connection.enqueue_upsert(str(self._type), _strip_transient_fields(value.info)))
        return value

    def delete(self, key: str) -> None:
        self._dict.pop(key, None)
        asyncio.create_task(self._connection.enqueue_delete(str(self._type), key))

    def next(self):
        return next(iter(self._dict.items()))

    def empty(self) -> bool:
        return not bool(self._dict)

    def has_downloads(self) -> bool:
        return any(self._dict[key].info.auto_start and self._dict[key].started() is False for key in self._dict)

    def get_next_download(self) -> Download | None:
        for key in self._dict:
            ref = self._dict[key]
            if ref.info.auto_start and not ref.started() and not ref.is_cancelled():
                return ref
        return None

    async def get_total_count(self, status_filter: str | None = None) -> int:
        return await self._connection.count(str(self._type), status_filter=status_filter)

    async def get_items_paginated(
        self, page: int = 1, per_page: int = 50, order: str = "DESC", status_filter: str | None = None
    ):
        if page < 1:
            msg = "page must be >= 1"
            raise ValueError(msg)
        if per_page < 1:
            msg = "per_page must be >= 1"
            raise ValueError(msg)
        order = order.upper()
        if order not in ("ASC", "DESC"):
            msg = f"order must be 'ASC' or 'DESC', got '{order}'"
            raise ValueError(msg)
        return await self._connection.paginate(str(self._type), page, per_page, order, status_filter)

    async def bulk_delete(self, ids: Iterable[str]) -> int:
        deleted = await self._connection.bulk_delete(str(self._type), ids)
        for _id in ids:
            self._dict.pop(_id, None)
        return deleted

    async def test(self) -> bool:
        await self._connection.count(str(self._type))
        return True