# Copyright 2013-2025 The py-lmdb authors, all rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted only as authorized by the OpenLDAP
# Public License.
#
# A copy of this license is available in the file LICENSE in the
# top-level directory of the distribution or, alternatively, at
# <http://www.OpenLDAP.org/license.html>.
#
# OpenLDAP is a registered trademark of the OpenLDAP Foundation.
#
# Individual files and/or contributed packages may be copyright by
# other parties and/or subject to additional restrictions.
#
# This work also contains materials derived from public sources.
#
# Additional information about OpenLDAP can be obtained at
# <http://www.openldap.org/>.

from pprint import pprint
import atexit
import gzip
import itertools
import os
import shutil
import sys
import tempfile

from time import time as now
import random
import zlmdb.lmdb as lmdb

MAP_SIZE = 1048576 * 400
DB_PATH = "/ram/testdb"
USE_SPARSE_FILES = sys.platform != "darwin"

if os.path.exists("/ram"):
    DB_PATH = "/ram/testdb"
else:
    DB_PATH = tempfile.mktemp(prefix="dirtybench")


env = None


@atexit.register
def cleanup():
    if env:
        env.close()
    if os.path.exists(DB_PATH):
        shutil.rmtree(DB_PATH)


def reopen_env(**kwargs):
    if env:
        env.close()
    if os.path.exists(DB_PATH):
        shutil.rmtree(DB_PATH)
    return lmdb.open(DB_PATH, map_size=MAP_SIZE, writemap=USE_SPARSE_FILES, **kwargs)


def case(title, **params):
    def wrapper(func):
        t0 = now()
        count = func()
        t1 = now()
        print("%40s:  %2.3fs   %8d/sec" % (title, t1 - t0, count / (t1 - t0)))
        return func

    return wrapper


def x():
    big = b""  # b'*' * 400

    t0 = now()
    words_path = os.path.join(os.path.dirname(__file__), "words.gz")
    words = set(gzip.open(words_path, "rt").read().splitlines())
    words.update([w.upper() for w in words])
    words.update([w[::-1] for w in words])
    words.update([w[::-1].upper() for w in words])
    words.update(["-".join(w) for w in words])
    # words.update(['+'.join(w) for w in words])
    # words.update(['/'.join(w) for w in words])
    words = [w.encode("utf-8") for w in words]
    alllen = sum(len(w) for w in words)
    avglen = alllen / len(words)
    print(
        "permutate %d words avglen %d took %.2fsec" % (len(words), avglen, now() - t0)
    )
    print("DB_PATH:", DB_PATH)

    words_sorted = sorted(words)
    items = [(w, big or w) for w in words]
    items_sorted = [(w, big or w) for w in words_sorted]

    global env
    env = reopen_env()

    @case("insert")
    def test():
        with env.begin(write=True) as txn:
            for word in words:
                txn.put(word, big or word)
            return len(words)

    st = env.stat()
    print
    print("stat:", st)
    print(
        "k+v size %.2fkb avg %d, on-disk size: %.2fkb avg %d"
        % (
            (2 * alllen) / 1024.0,
            (2 * alllen) / len(words),
            (st["psize"] * st["leaf_pages"]) / 1024.0,
            (st["psize"] * st["leaf_pages"]) / len(words),
        )
    )
    print()

    @case("enum (key, value) pairs")
    def test():
        with env.begin() as txn:
            return sum(1 for _ in txn.cursor())

    @case("reverse enum (key, value) pairs")
    def test():
        with env.begin() as txn:
            return sum(1 for _ in txn.cursor().iterprev())

    @case("enum (key, value) buffers")
    def test():
        with env.begin(buffers=True) as txn:
            return sum(1 for _ in txn.cursor())

    print

    @case("rand lookup")
    def test():
        with env.begin() as txn:
            for word in words:
                txn.get(word)
        return len(words)

    @case("per txn rand lookup")
    def test():
        for word in words:
            with env.begin() as txn:
                txn.get(word)
        return len(words)

    @case("rand lookup+hash")
    def test():
        with env.begin() as txn:
            for word in words:
                hash(txn.get(word))
        return len(words)

    @case("rand lookup buffers")
    def test():
        with env.begin(buffers=True) as txn:
            for word in words:
                txn.get(word)
        return len(words)

    @case("rand lookup+hash buffers")
    def test():
        with env.begin(buffers=True) as txn:
            for word in words:
                hash(bytes(txn.get(word)))
        return len(words)

    @case("rand lookup buffers (cursor)")
    def test():
        with env.begin(buffers=True) as txn:
            cursget = txn.cursor().get
            for word in words:
                cursget(word)
        return len(words)

    print

    @case("get+put")
    def test():
        with env.begin(write=True) as txn:
            for word in words:
                txn.get(word)
                txn.put(word, word)
            return len(words)

    @case("replace")
    def test():
        with env.begin(write=True) as txn:
            for word in words:
                txn.replace(word, word)
        return len(words)

    @case("get+put (cursor)")
    def test():
        with env.begin(write=True) as txn:
            with txn.cursor() as cursor:
                for word in words:
                    cursor.get(word)
                    cursor.put(word, word)
            return len(words)

    @case("replace (cursor)")
    def test():
        with env.begin(write=True) as txn:
            with txn.cursor() as cursor:
                for word in words:
                    cursor.replace(word, word)
        return len(words)

    print

    env = reopen_env()

    @case("insert (rand)")
    def test():
        with env.begin(write=True) as txn:
            for word in words:
                txn.put(word, big or word)
        return len(words)

    env = reopen_env()

    @case("insert (seq)")
    def test():
        with env.begin(write=True) as txn:
            for word in words_sorted:
                txn.put(word, big or word)
        return len(words)

    env = reopen_env()

    @case("insert (rand), reuse cursor")
    def test():
        with env.begin(write=True) as txn:
            curs = txn.cursor()
            for word in words:
                curs.put(word, big or word)
        return len(words)

    env = reopen_env()

    @case("insert (seq), reuse cursor")
    def test():
        with env.begin(write=True) as txn:
            curs = txn.cursor()
            for word in words_sorted:
                curs.put(word, big or word)
        return len(words)

    env = reopen_env()

    @case("insert, putmulti")
    def test():
        with env.begin(write=True) as txn:
            txn.cursor().putmulti(items)
        return len(words)

    env = reopen_env()

    @case("insert, putmulti+generator")
    def test():
        with env.begin(write=True) as txn:
            txn.cursor().putmulti((w, big or w) for w in words)
        return len(words)

    print

    env = reopen_env()

    @case("append")
    def test():
        with env.begin(write=True) as txn:
            for word in words_sorted:
                txn.put(word, big or word, append=True)
        return len(words)

    env = reopen_env()

    @case("append, reuse cursor")
    def test():
        with env.begin(write=True) as txn:
            curs = txn.cursor()
            for word in words_sorted:
                curs.put(word, big or word, append=True)
        return len(words)

    env = reopen_env()

    @case("append+putmulti")
    def test():
        with env.begin(write=True) as txn:
            txn.cursor().putmulti(items_sorted, append=True)
        return len(words)

    print
    st = env.stat()
    print("stat:", st)
    print(
        "k+v size %.2fkb avg %d, on-disk size: %.2fkb avg %d"
        % (
            (2 * alllen) / 1024.0,
            (2 * alllen) / len(words),
            (st["psize"] * st["leaf_pages"]) / 1024.0,
            (st["psize"] * st["leaf_pages"]) / len(words),
        )
    )


x()
