"""
Helper functions to access Archive and FollowTheMoney data within Jobs
"""

from contextlib import contextmanager
from pathlib import Path
from typing import ContextManager, Generator, Iterable

from anystore.logic.virtual import VirtualIO
from followthemoney import EntityProxy
from ftmq.store.fragments import get_fragments
from ftmq.store.fragments.loader import BulkLoader

from openaleph_procrastinate.archive import get_archive, lookup_key
from openaleph_procrastinate.exceptions import EntityNotFound
from openaleph_procrastinate.settings import OpenAlephSettings

OPAL_ORIGIN = "openaleph_procrastinate"
ARCHIVE_CHECKSUM_ALGORITHM = "sha1"
settings = OpenAlephSettings()
sqlalchemy_pool = {
    "pool_size": settings.db_pool_size,
    "max_overflow": settings.db_pool_size,
}


def get_localpath(dataset: str, content_hash: str) -> ContextManager[Path]:
    """
    Load a file from the archive and store it in a local temporary path for
    further processing. The file is cleaned up after leaving the context.
    [Reference][openaleph_procrastinate.model.DatasetJob.get_file_references]
    """
    archive = get_archive()
    key = lookup_key(content_hash)
    return archive.local_path(key)


def open_file(dataset: str, content_hash: str) -> ContextManager[VirtualIO]:
    """
    Load a file from the archive and store it in a local temporary path for
    further processing. Returns an open file handler. The file is closed and
    cleaned up after leaving the context.
    [Reference][openaleph_procrastinate.model.DatasetJob.get_file_references]
    """
    archive = get_archive()
    key = lookup_key(content_hash)
    return archive.local_open(key, algorithm=ARCHIVE_CHECKSUM_ALGORITHM)


def load_entity(dataset: str, entity_id: str) -> EntityProxy:
    """
    Retrieve a single entity from the store.
    """
    store = get_fragments(
        dataset, database_uri=settings.fragments_uri, **sqlalchemy_pool
    )
    entity = store.get(entity_id)
    if entity is None:
        raise EntityNotFound(f"Entity `{entity_id}` not found in dataset `{dataset}`")
    return entity


def load_entities(
    dataset: str, entity_ids: Iterable[str]
) -> Generator[EntityProxy, None, None]:
    """
    Batch retrieve entities from the fragment store.
    """
    store = get_fragments(
        dataset, database_uri=settings.fragments_uri, **sqlalchemy_pool
    )
    yield from store.iterate(entity_ids)


@contextmanager
def entity_writer(
    dataset: str, origin: str = OPAL_ORIGIN
) -> Generator[BulkLoader, None, None]:
    """
    Get the `ftmq.store.fragments.BulkLoader` for the given `dataset`. The
    writer is flushed when leaving the context.
    """
    store = get_fragments(
        dataset,
        origin=origin,
        database_uri=settings.fragments_uri,
        **sqlalchemy_pool,
    )
    loader = store.bulk()
    try:
        yield loader
    finally:
        loader.flush()
