#!/usr/bin/env python2
# -*- coding: utf-8 -*-
# Copyright (C) 2018-2021  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.
"""zblk_test_gen.py generates test data for zblk_test.go"""

from ZODB.DB import DB
from ZODB.utils import u64
from wendelin.bigfile.file_zodb import ZBlk0, ZBlk1, ZBigFile
from BTrees.IOBTree import IOBTree, IOBucket
from numpy import arange
import os, os.path, transaction

import zodbtools.test.gen_testdata  # to make time predictable
from zodbtools.test.gen_testdata import run_with_zodb4py2_compat

K = 1024

def main():
    run_with_zodb4py2_compat(main2)

def main2():
    outfs = "testdata/zblk.fs"
    rm_f(outfs)
    rm_f(outfs + ".index")
    db = DB(outfs)
    conn = db.open()
    root = conn.root()

    # one ZBigFile with two data blocks
    root['zbigf'] = zf = ZBigFile(2*1024*K)
    root['zblk0'] = z0 = ZBlk0()
    root['zblk1'] = z1 = ZBlk1()    # also covers ZData

    # zblk0 with small data
    z0.setblkdata(brange32(16*K))
    # zblk1 with head + large hole + tail.
    # head and tail are 128K because if smaller, z1.chunktab would serialize without buckets.
    z1ht = 128*K
    z1.setblkdata(brange32(z1ht) + bzeros(zf.blksize - 2*z1ht) + breverse(brange32(z1ht)))
    zf.blktab[1] = z0
    zf.blktab[3] = z1

    # make sure there is at least one bucket in z1's chunktab
    #
    # we need at least one bucket to verify how zblk.go handles chunktab bucket loading;
    # more than 1 bucket just wastes space, however with only 1 bucket a tree
    # is serialized to have leaf nodes directly in the tree state.
    # -> verify for ==2 buckets.
    assertIOBTreeHas2Buckets(z1.chunktab)

    transaction.commit()

    with open("ztestdata_zblk_test.go", "w") as f:
        def emit(v):
            print >>f, v
        emit("// Code generated by %s; DO NOT EDIT." % __file__)
        emit("package zdata\n")
        emit('import "lab.nexedi.com/kirr/neo/go/zodb"\n')

        emit("const zf_blksize = %d" % zf.blksize)
        emit("const zf_size    = %d" % ((zf.blktab.maxKey()+1)*zf.blksize))
        emit("const z0_oid     = zodb.Oid(%d)" % u64(z0._p_oid))
        emit("const z1_oid     = zodb.Oid(%d)" % u64(z1._p_oid))
        emit("const zf_oid     = zodb.Oid(%d)" % u64(zf._p_oid))
        emit("const z0_rev     = zodb.Tid(0x%x)"  % u64(z0._p_serial))
        emit("const z1_rev     = zodb.Tid(0x%x)"  % u64(z1._p_serial))
        emit("const z0_len     = %d" % len(z0.loadblkdata()))
        emit("const z1_htlen   = %d" % z1ht)


    conn.close()
    db.close()


# brange32 returns bytes with big-endian uint32 sequence filling them.
# returned bytes has len == size.
def brange32(size):
    # 0, 1, 2, ... as u32
    return arange(0, size//4, dtype='>u4').tobytes()

# bzeros returns bytes of requested size with 0 filling them.
def bzeros(size):
    return b'\0'*size

# breverse returns bytes in the reverse order.
def breverse(b):
    assert isinstance(b, bytes)
    _ = bytearray(b)
    _.reverse()
    return bytes(_)

# assertIOBTreeHas2Buckets asserts that IOBTree has 2 buckets.
def assertIOBTreeHas2Buckets(t):
    assert isinstance(t, IOBTree)

    # https://github.com/zopefoundation/BTrees/blob/4.5.0-1-gc8bf24e/BTrees/BTreeTemplate.c#L1087-L1109
    _ = t.__getstate__()
    assert len(_) == 2, (len(_), _)
    assert len(_[0]) == 3, _[0]                   # (bucket0, key1, bucket1)
    assert isinstance(_[0][0], IOBucket), _[0][0] # bucket0
    assert isinstance(_[0][1], int),      _[0][2] # key1
    assert isinstance(_[0][2], IOBucket), _[0][2] # bucket1
    assert isinstance(_[1],    IOBucket), _[1]    # .firstbucket

# rm_f is like `rm -f` in shell.
def rm_f(path):
    if os.path.exists(path):
        os.remove(path)


if __name__ == '__main__':
    main()
