# Copyright (C) 2020 The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information

from typing import Dict, Generator, Iterable

from swh.model.model import Content
from swh.objstorage import get_objstorage
from swh.objstorage.exc import ObjNotFoundError

from .exc import StorageArgumentException


class ObjStorage:
    """Objstorage collaborator in charge of adding objects to
    the objstorage.

    """
    def __init__(self, objstorage_config: Dict):
        self.objstorage = get_objstorage(**objstorage_config)

    def __getattr__(self, key):
        if key == 'objstorage':
            raise AttributeError(key)
        return getattr(self.objstorage, key)

    def content_get(self, contents: Iterable[bytes]) -> Generator:
        """Retrieve content data from the objstorage

        Args:
            contents: List of contents to retrieve data from

        """
        for obj_id in contents:
            try:
                data = self.objstorage.get(obj_id)
            except ObjNotFoundError:
                yield None
                continue

            yield {'sha1': obj_id, 'data': data}

    def content_add(self, contents: Iterable[Content]) -> Dict:
        """Add contents to the objstorage.

        Args:
            contents: List of contents to add1

        Returns:
            The summary dict of content and content bytes added to the
            objstorage.

        """
        contents = list(contents)
        if any(cont.data is None for cont in contents):
            raise StorageArgumentException('Missing data')
        summary = self.objstorage.add_batch({
            cont.sha1: cont.data
            for cont in contents
        })
        return {
            'content:add': summary['object:add'],
            'content:add:bytes': summary['object:add:bytes']
        }