# simplestorerepo.py - Extension that swaps in alternate repository storage.

#

# Copyright 2018 Gregory Szorc <gregory.szorc@gmail.com>

#

# This software may be used and distributed according to the terms of the

# GNU General Public License version 2 or any later version.

# To use this with the test suite:

#

#   $ HGREPOFEATURES="simplestore" ./run-tests.py \

#       --extra-config-opt extensions.simplestore=`pwd`/simplestorerepo.py

import stat

from typing import (

    Iterable,

    Iterator,

    Optional,

)

from mercurial.i18n import _

from mercurial.node import (

    bin,

    hex,

    nullrev,

)

from mercurial.thirdparty import attr

from mercurial import (

    ancestor,

    bundlerepo,

    error,

    extensions,

    localrepo,

    mdiff,

    pycompat,

    revlog,

    store,

    verify,

)

from mercurial.interfaces import (

    repository,

)

from mercurial.utils import (

    cborutil,

    storageutil,

)

from mercurial.revlogutils import flagutil

# Note for extension authors: ONLY specify testedwith = 'ships-with-hg-core' for

# extensions which SHIP WITH MERCURIAL. Non-mainline extensions should

# be specifying the version(s) of Mercurial they are tested with, or

# leave the attribute unspecified.

testedwith = b'ships-with-hg-core'

REQUIREMENT = b'testonly-simplestore'

def validatenode(node):

    if isinstance(node, int):

        raise ValueError('expected node; got int')

    if len(node) != 20:

        raise ValueError('expected 20 byte node')

def validaterev(rev):

    if not isinstance(rev, int):

        raise ValueError('expected int')

class simplestoreerror(error.StorageError):

    pass

@attr.s(slots=True)

class simplestorerevisiondelta(repository.irevisiondelta):

    node = attr.ib(type=bytes)

    p1node = attr.ib(type=bytes)

    p2node = attr.ib(type=bytes)

    basenode = attr.ib(type=bytes)

    flags = attr.ib(type=int)

    baserevisionsize = attr.ib(type=Optional[int])

    revision = attr.ib(type=Optional[bytes])

    delta = attr.ib(type=Optional[bytes])

    sidedata = attr.ib(type=Optional[bytes])

    protocol_flags = attr.ib(type=int)

    linknode = attr.ib(default=None, type=Optional[bytes])

@attr.s(frozen=True)

class simplefilestoreproblem(repository.iverifyproblem):

    warning = attr.ib(default=None, type=Optional[bytes])

    error = attr.ib(default=None, type=Optional[bytes])

    node = attr.ib(default=None, type=Optional[bytes])

class filestorage(repository.ifilestorage):

    """Implements storage for a tracked path.

    Data is stored in the VFS in a directory corresponding to the tracked

    path.

    Index data is stored in an ``index`` file using CBOR.

    Fulltext data is stored in files having names of the node.

    """

    _flagserrorclass = simplestoreerror

    def __init__(self, repo, svfs, path):

        self.nullid = repo.nullid

        self._repo = repo

        self._svfs = svfs

        self._path = path

        self._storepath = b'/'.join([b'data', path])

        self._indexpath = b'/'.join([self._storepath, b'index'])

        indexdata = self._svfs.tryread(self._indexpath)

        if indexdata:

            indexdata = cborutil.decodeall(indexdata)

        self._indexdata = indexdata or []

        self._indexbynode = {}

        self._indexbyrev = {}

        self._index = []

        self._refreshindex()

        self._flagprocessors = dict(flagutil.flagprocessors)

    def _refreshindex(self):

        self._indexbynode.clear()

        self._indexbyrev.clear()

        self._index = []

        for i, entry in enumerate(self._indexdata):

            self._indexbynode[entry[b'node']] = entry

            self._indexbyrev[i] = entry

        self._indexbynode[self._repo.nullid] = {

            b'node': self._repo.nullid,

            b'p1': self._repo.nullid,

            b'p2': self._repo.nullid,

            b'linkrev': nullrev,

            b'flags': 0,

        }

        self._indexbyrev[nullrev] = {

            b'node': self._repo.nullid,

            b'p1': self._repo.nullid,

            b'p2': self._repo.nullid,

            b'linkrev': nullrev,

            b'flags': 0,

        }

        for i, entry in enumerate(self._indexdata):

            p1rev, p2rev = self.parentrevs(self.rev(entry[b'node']))

            # start, length, rawsize, chainbase, linkrev, p1, p2, node

            self._index.append(

                (0, 0, 0, -1, entry[b'linkrev'], p1rev, p2rev, entry[b'node'])

            )

        self._index.append((0, 0, 0, -1, -1, -1, -1, self._repo.nullid))

    def __len__(self) -> int:

        return len(self._indexdata)

    def __iter__(self) -> Iterator[int]:

        return iter(range(len(self)))

    def revs(self, start=0, stop=None):

        step = 1

        if stop is not None:

            if start > stop:

                step = -1

            stop += step

        else:

            stop = len(self)

        return range(start, stop, step)

    def parents(self, node):

        validatenode(node)

        if node not in self._indexbynode:

            raise KeyError('unknown node')

        entry = self._indexbynode[node]

        return entry[b'p1'], entry[b'p2']

    def parentrevs(self, rev):

        p1, p2 = self.parents(self._indexbyrev[rev][b'node'])

        return self.rev(p1), self.rev(p2)

    def rev(self, node):

        validatenode(node)

        try:

            self._indexbynode[node]

        except KeyError:

            raise error.LookupError(node, self._indexpath, _('no node'))

        for rev, entry in self._indexbyrev.items():

            if entry[b'node'] == node:

                return rev

        raise error.ProgrammingError(b'this should not occur')

    def node(self, rev):

        validaterev(rev)

        return self._indexbyrev[rev][b'node']

    def hasnode(self, node):

        validatenode(node)

        return node in self._indexbynode

    def censorrevision(self, tr, censornode, tombstone=b''):

        raise NotImplementedError('TODO')

    def lookup(self, node):

        if isinstance(node, int):

            return self.node(node)

        if len(node) == 20:

            self.rev(node)

            return node

        try:

            rev = int(node)

            if '%d' % rev != node:

                raise ValueError

            if rev < 0:

                rev = len(self) + rev

            if rev < 0 or rev >= len(self):

                raise ValueError

            return self.node(rev)

        except (ValueError, OverflowError):

            pass

        if len(node) == 40:

            try:

                rawnode = bin(node)

                self.rev(rawnode)

                return rawnode

            except TypeError:

                pass

        raise error.LookupError(node, self._path, _('invalid lookup input'))

    def linkrev(self, rev):

        validaterev(rev)

        return self._indexbyrev[rev][b'linkrev']

    def _flags(self, rev):

        validaterev(rev)

        return self._indexbyrev[rev][b'flags']

    def _candelta(self, baserev, rev):

        validaterev(baserev)

        validaterev(rev)

        if (self._flags(baserev) & revlog.REVIDX_RAWTEXT_CHANGING_FLAGS) or (

            self._flags(rev) & revlog.REVIDX_RAWTEXT_CHANGING_FLAGS

        ):

            return False

        return True

    def checkhash(self, text, node, p1=None, p2=None, rev=None):

        if p1 is None and p2 is None:

            p1, p2 = self.parents(node)

        if node != storageutil.hashrevisionsha1(text, p1, p2):

            raise simplestoreerror(

                _("integrity check failed on %s") % self._path

            )

    def revision(self, nodeorrev, raw=False):

        if isinstance(nodeorrev, int):

            node = self.node(nodeorrev)

        else:

            node = nodeorrev

        validatenode(node)

        if node == self._repo.nullid:

            return b''

        rev = self.rev(node)

        flags = self._flags(rev)

        path = b'/'.join([self._storepath, hex(node)])

        rawtext = self._svfs.read(path)

        if raw:

            validatehash = flagutil.processflagsraw(self, rawtext, flags)

            text = rawtext

        else:

            r = flagutil.processflagsread(self, rawtext, flags)

            text, validatehash = r

        if validatehash:

            self.checkhash(text, node, rev=rev)

        return text

    def rawdata(self, nodeorrev):

        return self.revision(raw=True)

    def read(self, node):

        validatenode(node)

        revision = self.revision(node)

        if not revision.startswith(b'\1\n'):

            return revision

        start = revision.index(b'\1\n', 2)

        return revision[start + 2 :]

    def renamed(self, node):

        validatenode(node)

        if self.parents(node)[0] != self._repo.nullid:

            return False

        fulltext = self.revision(node)

        m = storageutil.parsemeta(fulltext)[0]

        if m and 'copy' in m:

            return m['copy'], bin(m['copyrev'])

        return False

    def cmp(self, node, text):

        validatenode(node)

        t = text

        if text.startswith(b'\1\n'):

            t = b'\1\n\1\n' + text

        p1, p2 = self.parents(node)

        if storageutil.hashrevisionsha1(t, p1, p2) == node:

            return False

        if self.iscensored(self.rev(node)):

            return text != b''

        if self.renamed(node):

            t2 = self.read(node)

            return t2 != text

        return True

    def size(self, rev):

        validaterev(rev)

        node = self._indexbyrev[rev][b'node']

        if self.renamed(node):

            return len(self.read(node))

        if self.iscensored(rev):

            return 0

        return len(self.revision(node))

    def iscensored(self, rev):

        validaterev(rev)

        return self._flags(rev) & repository.REVISION_FLAG_CENSORED

    def commonancestorsheads(self, a, b):

        validatenode(a)

        validatenode(b)

        a = self.rev(a)

        b = self.rev(b)

        ancestors = ancestor.commonancestorsheads(self.parentrevs, a, b)

        return pycompat.maplist(self.node, ancestors)

    def descendants(self, revs):

        # This is a copy of revlog.descendants()

        first = min(revs)

        if first == nullrev:

            for i in self:

                yield i

            return

        seen = set(revs)

        for i in self.revs(start=first + 1):

            for x in self.parentrevs(i):

                if x != nullrev and x in seen:

                    seen.add(i)

                    yield i

                    break

    # Required by verify.

    def files(self):

        entries = self._svfs.listdir(self._storepath)

        # Strip out undo.backup.* files created as part of transaction

        # recording.

        entries = [f for f in entries if not f.startswith('undo.backup.')]

        return [b'/'.join((self._storepath, f)) for f in entries]

    def storageinfo(

        self,

        exclusivefiles=False,

        sharedfiles=False,

        revisionscount=False,

        trackedsize=False,

        storedsize=False,

    ):

        # TODO do a real implementation of this

        return {

            'exclusivefiles': [],

            'sharedfiles': [],

            'revisionscount': len(self),

            'trackedsize': 0,

            'storedsize': None,

        }

    def verifyintegrity(self, state) -> Iterable[repository.iverifyproblem]:

        state['skipread'] = set()

        for rev in self:

            node = self.node(rev)

            try:

                self.revision(node)

            except Exception as e:

                yield simplefilestoreproblem(

                    error='unpacking %s: %s' % (node, e), node=node

                )

                state['skipread'].add(node)

    def emitrevisions(

        self,

        nodes,

        nodesorder=None,

        revisiondata=False,

        assumehaveparentrevisions=False,

        deltamode=repository.CG_DELTAMODE_STD,

        sidedata_helpers=None,

    ):

        # TODO this will probably break on some ordering options.

        nodes = [n for n in nodes if n != self._repo.nullid]

        if not nodes:

            return

        yield from storageutil.emitrevisions(

            self,

            nodes,

            nodesorder,

            simplestorerevisiondelta,

            revisiondata=revisiondata,

            assumehaveparentrevisions=assumehaveparentrevisions,

            deltamode=deltamode,

            sidedata_helpers=sidedata_helpers,

        )

    def add(self, text, meta, transaction, linkrev, p1, p2):

        if meta or text.startswith(b'\1\n'):

            text = storageutil.packmeta(meta, text)

        return self.addrevision(text, transaction, linkrev, p1, p2)

    def addrevision(

        self,

        text,

        transaction,

        linkrev,

        p1,

        p2,

        node=None,

        flags=revlog.REVIDX_DEFAULT_FLAGS,

        cachedelta=None,

    ):

        validatenode(p1)

        validatenode(p2)

        if flags:

            node = node or storageutil.hashrevisionsha1(text, p1, p2)

        rawtext, validatehash = flagutil.processflagswrite(self, text, flags)

        node = node or storageutil.hashrevisionsha1(text, p1, p2)

        if node in self._indexbynode:

            return node

        if validatehash:

            self.checkhash(rawtext, node, p1=p1, p2=p2)

        return self._addrawrevision(

            node, rawtext, transaction, linkrev, p1, p2, flags

        )