"""Bluesky repo storage base class and in-memory implementation.

Lightly based on:
https://github.com/bluesky-social/atproto/blob/main/packages/repo/src/storage/repo-storage.ts
"""
from collections import namedtuple

import dag_cbor.decoding, dag_cbor.encoding
from multiformats import CID, multicodec, multihash

from .util import dag_cbor_cid


CommitData = namedtuple('CommitData', [
  'commit',  # CID
  'blocks',  # BlockMap
  'prev',    # CID or None
])


class BlockMap(dict):
    """dict subclass that stores blocks as CID => blocks (bytes) mappings.

    A block is a DAG-CBOR encoded node, ie a record, MST entry, or commit.
    """
    def add(self, val):
        """Encodes a value as a block and adds it.

        TODO: remove or refactor? keep the invariant that BlockMap stores
        blocks, Storage stores records?

        Args:
          val: dict, record

        Returns:
          :class:`CID`
        """
        block = dag_cbor.encoding.encode(val)
        digest = multihash.digest(block, 'sha2-256')
        cid = CID('base58btc', 1, multicodec.get('dag-cbor'), digest)
        self[cid] = block
        return cid

    def byte_size(self):
        """Returns the cumulative size of all blocks, in bytes.

        Returns:
          int
        """
        return sum(len(b) for b in self.values())


class Storage:
    """Abstract base class for storing nodes: records, MST entries, and commits.

    Concrete subclasses should implement this on top of physical storage,
    eg database, filesystem, in memory.

    # TODO: batch operations?

    Attributes:
      head: :class:`CID`
    """
    head = None

    def read(self, cid):
        """Reads a node from storage.

        Args:
          cid: :class:`CID`

        Returns:
          dict, a record, commit, or serialized MST node
        """
        raise NotImplementedError()

    def read_many(self, cids):
        """Batch read multiple nodes from storage.

        Args:
          sequence of :class:`CID`

        Returns:
          tuple: (dict {:class:`CID`: dict node},
                  sequence of :class:`CID` that weren't found)
        """
        raise NotImplementedError()

    def read_blocks(self, cids):
        """Batch read multiple blocks from storage.

        Args:
          sequence of :class:`CID`

        Returns:
          tuple: (:class:`BlockMap` with found blocks,
                  sequence of :class:`CID` that weren't found)
        """
        raise NotImplementedError()

    def has(self, cid):
        """Checks if a given :class:`CID` is currently stored.

        Args:
          cid: :class:`CID`

        Returns:
          boolean
        """
        raise NotImplementedError()

    def write(self, node):
        """Writes a node to storage.

        Args:
          node: a record, commit, or serialized MST node

        Returns:
          :class:`CID`
        """
        raise NotImplementedError()

    def apply_commit(self, commit_data):
        """Writes a commit to storage.

        Args:
          commit: :class:`CommitData`
        """
        raise NotImplementedError()


class MemoryStorage(Storage):
    """In memory storage implementation.

    Attributes:
      blocks: :class:`BlockMap`
    """
    blocks = None

    def __init__(self):
        self.blocks = BlockMap()

    def read(self, cid):
        return dag_cbor.decoding.decode(self.blocks[cid])

    def read_many(self, cids):
        blocks, missing = self.read_blocks(cids)
        nodes = {cid: dag_cbor.decoding.decode(block)
                 for cid, block in blocks.items()}
        return nodes, missing

    def read_blocks(self, cids):
        found = {}
        missing = []

        for cid in cids:
            block = self.blocks.get(cid)
            if block:
                found[cid] = block
            else:
                missing.append(cid)

        return found, missing

    def has(self, cid):
        return cid in self.blocks

    def write(self, node):
        self.blocks.add(node)

    def apply_commit(self, commit_data):
        self.blocks.update(commit_data.blocks)
        self.head = commit_data.commit
