#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright (C) 2020  Nexedi SA and Contributors.
#                     Kirill Smelkov <kirr@nexedi.com>
#
# This program is free software: you can Use, Study, Modify and Redistribute
# it under the terms of the GNU General Public License version 3, or (at your
# option) any later version, as published by the Free Software Foundation.
#
# You can also Link and Combine this program with other software covered by
# the terms of any of the Free Software licenses or any of the Open Source
# Initiative approved licenses and Convey the resulting work. Corresponding
# source of such a combination shall include the source code for all other
# software used.
#
# This program is distributed WITHOUT ANY WARRANTY; without even the implied
# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#
# See COPYING file for full licensing terms.
# See https://www.nexedi.com/licensing for rationale and options.
"""Program treegen provides infrastructure to generate ZODB BTree states.

It is used as helper for ΔBTree tests.

The following subcommands are provided:

- `trees` transition ZODB tree through requested tree states,
- `allstructs` generates topologies for subset of all possible tree changes in
  between two trees specified by two key->value dicts.

Because python/pkg_resources startup is very slow(*) all subcommands can be
used either in CLI or in server mode, where requests are continuously read from
stdin.

Subcommands documentation follows:


trees
-----

`treegen trees <zurl>` transitions ZODB LOBTree through requested tree states.
Tree states are specified on stdin as topology-encoded strings, 1 state per 1 line.
For every request the tree is changed to have specified keys, values and
topology, and modifications are committed to database. For every made commit
corresponding transaction ID is printed to stdout.

The tree, that `treegen trees` generates and works on, is accessible via
zconn.root()['treegen/tree'].

Trees protocol specification:

    S: tree.srv start @<head> root=<tree-root-oid>
    C: <tree>
    S: <tid>
    C: <tree>
    S: <tid>
    ...

session example:

    S: tree.srv start @03d85dd71ed0d2ee root=000000000000000b
    C: T/B1:a
    S: 03d85dd84fed7844
    C: T2/B1:a-B3:c
    S: 03d85dd871718899
    ...


allstructs
----------

`treegen allstructs` generates topologies for subset of all possible tree
changes in between two trees specified by two key->value dicts.

For every kv the following tree topologies are considered: 1) native (the one
that ZODB would usually create natively via regular usage), and 2) n random
ones. Then tree topologies are emitted corresponding to tree1->tree2 and
tree1<-tree2 transitions for all combinations of (tree1, tree2) pairs.

The output of `treegen allstructs` is valid input for `treegen trees`.

Allstructs protocol specification:

    S: # allstructs.srv start
    C: <maxdepth> <maxsplit> <n>(/<seed>) <kv1> <kv2>
    S: # allstructs <kv1> <kv2>
    S: # maxdepth=<maxdepth> maxsplit=<maxsplit> n=<n> seed=<seed>
    S: <tree1₀>
    S: <tree2₀>
    S: <tree1₁>
    ...
    S: # ----

session example:

    # allstructs.srv start
    1 1 10 1:a 2:b
    # allstructs 1:a 2:b
    # maxdepth=1 maxsplit=1 n=10 seed=1591369961
    T/B1:a
    T/B2:b
    T/T/B1:a
    T/B2:b
    T/B1:a
    T/T/B2:b
    T/T/B1:a
    T/T/B2:b
    T/B1:a
    # ----

--------

(*) 300-500ms, see https://github.com/pypa/setuptools/issues/510.
"""

from __future__ import print_function, absolute_import

import sys
from golang import func, defer, panic
from golang import time
from ZODB import DB
from ZODB.Connection import Connection
from ZODB.MappingStorage import MappingStorage
import transaction
import random
import six

from wendelin.wcfs.internal import xbtree, xbtree_test
from wendelin.bigfile.file_zodb import ZBlk
from zodbtools.util import storageFromURL, ashex

from persistent import CHANGED
from persistent.mapping import PersistentMapping

# XXX hack: set LOBTree.LOBTree -> XLOTree so that nodes are split often for
# regular tree updates (XLOTree is LOBTree with small .max_*_size). Do it this
# way so that generated database looks as if regular LOBTree was used. We use
# the hack because we cannot tune LOBTree directly.
XLOTree = xbtree_test.XLOTree
XLOTree.__module__ = 'BTrees.LOBTree'
XLOTree.__name__   = 'LOBTree'
import BTrees.LOBTree
BTrees.LOBTree.LOBTree = XLOTree
from BTrees.LOBTree import LOBTree


# ZCtx represent treegen-level connection to ZODB.
# It wraps zconn + provides treegen-specif integration.
class ZCtx(object):
    # .zconn
    # .root
    # .valdict = zconn.root['treegen/values'] = {} v -> ZBlk(v)

    # ZCtx(zstor) opens connection to zstor and initializes .valdict.
    def __init__(zctx, zstor):
        zctx.db    = DB(zstor)
        zctx.zconn = zctx.db.open()
        zctx.root  = zctx.zconn.root()

        # init root['treegen/values'] = {} v -> ZBlk(v)
        valdict = zctx.root.get('treegen/values', None)
        if valdict is None:
            valdict = zctx.root['treegen/values'] = PersistentMapping()
        valv = b'abcdefghij'
        for v in valv:
            zblk = valdict.get(v, None)
            if zblk is not None and zblk.loadblkdata() == v:
                continue
            zblk = ZBlk()
            zblk.setblkdata(v)
            valdict[v] = zblk
        zctx.valdict = valdict
        commit('treegen/values: init %r' % valv, skipIfEmpty=True)

    # close release resources associated with zctx.
    def close(zctx):
        zctx.zconn.close()
        zctx.db.close()

    # vdecode(vtxt) -> vobj  decodes value text into value object, e.g. 'a' -> ZBlk(a)
    # vencode(vobj) -> vtxt  encodes value object into value text, e.g. ZBlk(a) -> 'a'
    def vdecode(zctx, vtxt): # -> vobj
        return zctx.valdict[vtxt]
    def vencode(zctx, vobj): # -> vtxt
        for (k,v) in zctx.valdict.items():
            if v is vobj:
                return k
        raise KeyError("%r not found in value registry" % (vobj,))


# TreesSrv transitions ZODB tree through requested tree states.
# See top-level documentation for details.
@func
def TreesSrv(zstor, r):
    zctx = ZCtx(zstor)
    defer(zctx.close)

    ztree = zctx.root['treegen/tree'] = LOBTree()
    head = commit('treegen/tree: init')
    xprint("tree.srv start @%s root=%s" % (ashex(head), ashex(ztree._p_oid)))
    treetxtPrev = zctx.ztreetxt(ztree)

    for treetxt in xreadlines(r):
        # mark tree as changed if the same topology is requested twice.
        # this ensures we can actually make a non-empty commit
        if treetxt == treetxtPrev:
            ztree._p_changed = True
        treetxtPrev = treetxt

        tree = zctx.TopoDecode(treetxt)

        # treekv[k]=v for all k
        # do tree.keys() via walkBFS
        treekv = {}
        for level in xbtree._walkBFS(tree):
            for node in level:
                if isinstance(node, xbtree.Bucket):
                    if node.valuev is None:
                        panic("%s: tree must be {key->value}, not set" % treetxt)
                    assert len(node.keyv) == len(node.valuev)
                    for (k,v) in zip(node.keyv, node.valuev):
                        treekv[k] = v

        # change ztree to requested kv
        d = diff(ztree, treekv)
        patch(ztree, d, treekv)

        # restructure ztree to requested topology
        xbtree.Restructure(ztree, tree)

        # commit tree to storage
        head = commit("treegen/tree: %s" % treetxt)

        # verify what was persisted to storage is indeed what we wanted to persist
        zctx.zconn.cacheMinimize()
        treetxt_onstor = zctx.ztreetxt(ztree)
        if treetxt_onstor != treetxt:
            panic("BUG: tree wrongly saved to storage:\nsaved:     %s\nrequested: %s" %
                    (treetxt_onstor, treetxt))
        # verify saved tree consistency
        xbtree.zcheck(ztree)

        # ok
        xprint("%s" % ashex(head))


# AllStructsSrv is server version of AllStructs.
@func
def AllStructsSrv(r):
    xprint('# allstructs.srv start')
    for req in xreadlines(r):
        # maxdepth maxsplit n(/seed) kv1 kv2
        maxdepth, maxsplit, n, kv1txt, kv2txt = req.split()
        maxdepth = int(maxdepth)
        maxsplit = int(maxsplit)
        seed = None
        if '/' in n:
            n, seeds = n.split('/')
            seed = int(seeds)
        n = int(n)
        if kv1txt == 'ø': kv1txt = ''
        if kv2txt == 'ø': kv2txt = ''

        AllStructs(kv1txt, kv2txt, maxdepth, maxsplit, n, seed)
        xprint('# ----')

# AllStructs generates topologies for subset of all possible tree changes in
# between kv1 and kv2. See top-level documentation for details.
@func
def AllStructs(kv1txt, kv2txt, maxdepth, maxsplit, n, seed=None):
    zstor = MappingStorage() # in RAM storage to create native ZODB topologies
    zctx = ZCtx(zstor)
    defer(zctx.close)

    kv1 = kvDecode(kv1txt, zctx.vdecode)
    kv2 = kvDecode(kv2txt, zctx.vdecode)

    print("# allstructs %s %s" % (kv1txt, kv2txt))

    # create the tree
    ztree = zctx.root['ztree'] = LOBTree()
    commit('init')

    # initial kv1 and kv2 states with topologies prepared as ZODB would do natively
    patch(ztree, diff({}, kv1), verify=kv1)
    if kv1 == {}: ztree._p_changed = True # to avoid empty commit - see TreesSrv
    commit('kv1')
    t1struct0 = xbtree.StructureOf(ztree)

    patch(ztree, diff(kv1, kv2), verify=kv2)
    if kv2 == kv1: ztree._p_changed = True
    commit('kv2')
    t2struct0 = xbtree.StructureOf(ztree)

    # all tree topologies that can represent kv1 and kv2
    t1AllStructs = list(xbtree.AllStructs(kv1.keys(), maxdepth, maxsplit, kv=kv1))
    t2AllStructs = list(xbtree.AllStructs(kv2.keys(), maxdepth, maxsplit, kv=kv2))

    # seed
    if seed is None:
        seed = time.now()
    seed = int(seed)
    random.seed(seed)
    print("# maxdepth=%d maxsplit=%d n=%d seed=%d" % (maxdepth, maxsplit, n, seed))

    # all tree1 and tree2 topologies jumps in between we are going to emit:
    # native + n random ones.
    if t1struct0 in t1AllStructs: t1AllStructs.remove(t1struct0) # avoid dups
    if t2struct0 in t2AllStructs: t2AllStructs.remove(t2struct0)
    t1structv = [t1struct0] + random.sample(t1AllStructs, min(n, len(t1AllStructs)))
    t2structv = [t2struct0] + random.sample(t2AllStructs, min(n, len(t2AllStructs)))

    # emit topologies for tree1->tree2 and tree1<-tree2 transitions for all
    # combinations of tree1 and tree2.
    t12travel = list(bitravel2Way(t1structv, t2structv))
    for i,tstruct in enumerate(t12travel):
        if i%2 == 0:
            assert tstruct in t1structv
        else:
            assert tstruct in t2structv

        print(zctx.TopoEncode(tstruct))


# bitravel2Way generates travel path through all A<->B edges such
# that all edges a->b and a<-b are traveled and exactly once.
#
# The travel starts from A[0].
def bitravel2Way(A, B): # -> i[] of node
    na = len(A); assert na > 0
    nb = len(B); assert nb > 0

    yield A[0]                  # A₀
    for j in range(nb):
        yield B[j]              # A₀ -> Bj
        for i in range(1,na):
            yield A[i]          # Ai <- Bj
            yield B[j]          # Ai -> Bj
        yield A[0]              # A₀ <- Bj

def test_bitravel2Way():
    a,b,c = 'a','b','c'
    A = [a,b,c]
    B = [1, 2]
    got  = list(bitravel2Way(A, B))
    want = [a,1,b,1,c,1,a,2,b,2,c,2,a]
    assert got == want, (got, want)
test_bitravel2Way()


# kvEncode encodes key->value mapping into text.
# e.g. {1:'a', 2:'b'} -> '1:a,2:b'
def kvEncode(kvDict, vencode): # -> kvText
    retv = []
    for k in sorted(kvDict.keys()):
        v = kvDict[k]
        retv.append('%d:%s' % (k, vencode(v)))
    return ','.join(retv)

# kvDecode decodes key->value mapping from text.
# e.g. '1:a,2:b' -> {1:'a', 2:'b'}
def kvDecode(kvText, vdecode): # -> kvDict
    if kvText == "":
        return {}
    kv = {}
    for item in kvText.split(','):
        ktxt, vtxt = item.split(':')
        k = int(ktxt)
        v = vdecode(vtxt)
        if k in kv:
            raise ValueError("key %s present multiple times" % k)
        kv[k] = v
    return kv


# diff computes difference in between mappings d1 and d2.
DEL = 'ø'
def diff(d1, d2): # -> [] of (k,v) to change; DEL means del[k]
    delta = []
    keys = set(d1.keys())
    keys.update(d2.keys())
    for k in sorted(keys):
        v1 = d1.get(k, DEL)
        v2 = d2.get(k, DEL)
        if v1 is not v2:
            delta.append((k,v2))
    return delta

# patch changes mapping d according to diff.
# diff = [] of (k,v) to change; DEL means del[k]
def patch(d, diff, verify):
    for (k,v) in diff:
        if v is DEL:
            del d[k]
        else:
            d[k] = v

    if verify is None:
        return
    keys  = set(d.keys())
    keyok = set(verify.keys())
    if keys != keyok:
        panic("patch: verify: different keys: %s" % keys.symmetric_difference(keyok))
    for k in keys:
        if d[k] is not verify[k]:
            panic("patch: verify: [%d] different: got %r;  want %r" % (k, d[k], verify[k]))


# commit commits current transaction with description.
def commit(description, skipIfEmpty=False): # -> tid | None
    txn = transaction.get()
    if skipIfEmpty and len(txn._resources) == 0:
        return None
    txn.description = description
    # XXX hack to retrieve committed transaction ID via ._p_serial of object changed in this transaction
    assert len(txn._resources) == 1, txn._resources
    zconn = txn._resources[0]
    assert isinstance(zconn, Connection)
    assert len(zconn._registered_objects) > 0
    obj = zconn._registered_objects[0]
    # FIXME vvv fails sometims with obj._p_state=0 (UPTODATE)
    # see e.g. https://erp5.nexedi.net/test_result_module/20201016-89AA413F/5
    assert obj._p_state == CHANGED, obj._p_state
    txn.commit()
    return obj._p_serial


# ztreetxt returns text representation of a ZODB tree.
@func(ZCtx)
def ztreetxt(zctx, ztree): # -> txt
    assert ztree._p_jar is zctx.zconn
    return zctx.TopoEncode(xbtree.StructureOf(ztree))

# TopoEncode and TopoDecode perform trees encode/decode with values registry
# taken from zctx.valdict.
@func(ZCtx)
def TopoEncode(zctx, tree):
    return xbtree.TopoEncode(tree, zctx.vencode)
@func(ZCtx)
def TopoDecode(zctx, text):
    return xbtree.TopoDecode(text, zctx.vdecode)


# xreadlines iterates through lines in r skipping comments.
def xreadlines(r):
    while 1:
        l = r.readline()
        if l == '':
            break # EOF
        l = l.rstrip()  # trim trailing \n
        if l.startswith('#'):
            continue # skip comments
        yield l


@func
def cmd_allstructs(argv):
    if argv == ['-h']:
        print("Usage: cat requests |treegen allstructs",    file=sys.stderr)
        print("       treegen allstructs <request>",        file=sys.stderr)
        sys.exit(0)

    r = sys.stdin
    if len(argv) != 0:
        r = six.StringIO(' '.join(argv))

    AllStructsSrv(r)

@func
def cmd_trees(argv):
    if argv in ([], ['-h']):
        print("Usage: cat trees |treegen trees <zurl>",     file=sys.stderr)
        print("       treegen trees <zurl> <tree>+",        file=sys.stderr)
        sys.exit(1 if argv==[] else 0)

    zurl  = argv[0]
    zstor = storageFromURL(zurl)
    defer(zstor.close)

    r = sys.stdin
    treev = argv[1:]
    if len(treev) != 0:
        r = six.StringIO('\n'.join(treev))

    TreesSrv(zstor, r)


# xprint prints msg to stdout and flushes it.
def xprint(msg):
    print(msg)
    sys.stdout.flush()


cmdRegistry = {
    'allstructs':   cmd_allstructs,
    'trees':        cmd_trees,
}

def main():
    if len(sys.argv) < 2:
        print("Usage: treegen <command> ...", file=sys.stderr)
        sys.exit(1)

    cmd  = cmdRegistry.get(sys.argv[1], None)
    if cmd is None:
        print("E: treegen: unknown command %r" % sys.argv[1])
        sys.exit(1)
    argv = sys.argv[2:]
    cmd(argv)


if __name__ == '__main__':
    main()
