package com.webank.wecross.storage;

import java.util.HashSet;
import java.util.Set;
import java.util.function.BiConsumer;
import org.rocksdb.RocksDB;
import org.rocksdb.RocksDBException;
import org.rocksdb.WriteBatch;
import org.rocksdb.WriteOptions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class RocksDBBlockHeaderStorage implements BlockHeaderStorage {
    private boolean dbClosed = false;

    private RocksDB rocksDB;

    private static final String numberKey = "number";
    private static final String blockKeyPrefix = "block_";

    private Logger logger = LoggerFactory.getLogger(RocksDBBlockHeaderStorage.class);

    private Set<BiConsumer<Long, byte[]>> onBlockHeaderCallbacks = new HashSet<>();

    @Override
    public long readBlockNumber() {
        if (dbClosed == true) {
            logger.warn("Read RocksDB error: RocksDB has been closed");
            return -1;
        }

        try {
            byte[] blockNumberBytes = rocksDB.get(numberKey.getBytes());
            if (blockNumberBytes != null) {
                String blockNumberStr = new String(blockNumberBytes);
                long blockNumber = Long.valueOf(blockNumberStr);
                return blockNumber;
            } else {
                return -1;
            }
        } catch (RocksDBException e) {
            logger.error("Read RocksDB error", e);
        }
        return -1;
    }

    @Override
    public byte[] readBlockHeader(long blockNumber) {
        if (dbClosed == true) {
            logger.warn("Read RocksDB error: RocksDB has been closed");
            return null;
        }

        String key = "block_" + String.valueOf(blockNumber);

        try {
            return rocksDB.get(key.getBytes());
        } catch (RocksDBException e) {
            logger.error("RocksDB read error", e);
        }
        return null;
    }

    @Override
    public void writeBlockHeader(long blockNumber, byte[] blockHeader) {
        if (dbClosed == true) {
            logger.warn("Write RocksDB error: RocksDB has been closed");
            return;
        }

        String key = blockKeyPrefix + String.valueOf(blockNumber);

        try {
            WriteBatch writeBatch = new WriteBatch();
            writeBatch.put(numberKey.getBytes(), String.valueOf(blockNumber).getBytes());
            writeBatch.put(key.getBytes(), blockHeader);

            WriteOptions writeOptions = new WriteOptions();

            rocksDB.write(writeOptions, writeBatch);
            onBlockHeader(blockNumber, blockHeader);
        } catch (RocksDBException e) {
            logger.error("RocksDB write error", e);
        }
    }

    public RocksDB getRocksDB() {
        return rocksDB;
    }

    public void setRocksDB(RocksDB rocksDB) {
        this.rocksDB = rocksDB;
        this.dbClosed = false;
    }

    @Override
    public void close() {
        dbClosed = true;
        rocksDB.close();
    }

    @Override
    public void registerOnBlockHeader(BiConsumer<Long, byte[]> fn) {
        onBlockHeaderCallbacks.add(fn);
    }

    private void onBlockHeader(long blockNumber, byte[] blockHeader) {
        for (BiConsumer<Long, byte[]> fn : onBlockHeaderCallbacks) {
            fn.accept(blockNumber, blockHeader);
        }
    }
}