content: /*************************************************************************
  *                                                                        *
  *  This file is part of the 20n/act project.                             *
  *  20n/act enables DNA prediction for synthetic biology/bioengineering.  *
  *  Copyright (C) 2017 20n Labs, Inc.                                     *
  *                                                                        *
  *  Please direct all queries to act@20n.com.                             *
  *                                                                        *
  *  This program is free software: you can redistribute it and/or modify  *
  *  it under the terms of the GNU General Public License as published by  *
  *  the Free Software Foundation, either version 3 of the License, or     *
  *  (at your option) any later version.                                   *
  *                                                                        *
  *  This program is distributed in the hope that it will be useful,       *
  *  but WITHOUT ANY WARRANTY; without even the implied warranty of        *
  *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
  *  GNU General Public License for more details.                          *
  *                                                                        *
  *  You should have received a copy of the GNU General Public License     *
  *  along with this program.  If not, see <http://www.gnu.org/licenses/>. *
  *                                                                        *
  *************************************************************************/
  
  package com.act.utils;
  
  import com.act.utils.rocksdb.ColumnFamilyEnumeration;
  import com.act.utils.rocksdb.RocksDBAndHandles;
  import org.rocksdb.RocksDBException;
  
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  
  public class MockRocksDBAndHandles<T extends ColumnFamilyEnumeration<T>> extends RocksDBAndHandles<T> {
    T[] columnFamilies;
    Map<T, Map<List<Byte>, byte[]>> fakeDB = new HashMap<>();
  
    public MockRocksDBAndHandles(T[] columnFamilies) {
      super();
      this.columnFamilies = columnFamilies;
      for (T family : columnFamilies) {
        fakeDB.put(family, new HashMap<>());
      }
    }
  
    public static List<Byte> byteArrayToList(byte[] array) {
      // Must wrap byte[] for keys, as byte arrays hash code is address (and not content) based.
      List<Byte> keyList = new ArrayList<>(array.length);
      for (byte b : array) {
        keyList.add(b); // Streams and Arrays aren't friendly to byte[], so do it the old fashioned way.
      }
      return keyList;
    }
  
    public static byte[] byteListToArray(List<Byte> list) {
      byte[] array = new byte[list.size()];
      for (int i = 0; i < array.length; i++) {
        array[i] = list.get(i); // Sigh, still the old way with byte[].
      }
      return array;
    }
  
  
    public Map<T, Map<List<Byte>, byte[]>> getFakeDB() {
      return this.fakeDB;
    }
  
    @Override
    public void put(T columnFamily, byte[] key, byte[] val) throws RocksDBException {
      // Copy the value bytes since they'll likely be modified in place.
      fakeDB.get(columnFamily).put(byteArrayToList(key), Arrays.copyOf(val, val.length));
    }
  
    @Override
    public boolean keyMayExist(T columnFamily, byte[] key) throws RocksDBException {
      return fakeDB.get(columnFamily).containsKey(byteArrayToList(key));
    }
  
    @Override
    public byte[] get(T columnFamily, byte[] key) throws RocksDBException {
      return fakeDB.get(columnFamily).get(byteArrayToList(key));
    }
  
    @Override
    public RocksDBIterator newIterator(T columnFamily) throws RocksDBException {
      return new MockRocksDBIterator(fakeDB.get(columnFamily));
    }
  
    @Override
    public void flush(boolean waitForFlush) throws RocksDBException {
  
    }
  
    @Override
    public RocksDBWriteBatch<T> makeWriteBatch() {
      return new MockRocksDBWriteBatch<T>(this, 0);
    }
  
    public static class MockRocksDBWriteBatch<T extends ColumnFamilyEnumeration<T>>
        extends RocksDBWriteBatch<T> {
      RocksDBAndHandles<T> parent;
      Map<T, Map<List<Byte>, byte[]>> batch = new HashMap<>();
  
      protected MockRocksDBWriteBatch(RocksDBAndHandles<T> parent, int reservedBytes) {
        super();
        this.parent = parent;
      }
  
      @Override
      public void put(T columnFamily, byte[] key, byte[] val) throws RocksDBException {
        // This is not what a write batch is supposed to do, but it should get us close enough to test.
        Map<List<Byte>, byte[]> columnFamilyMap = batch.get(columnFamily);
        if (columnFamilyMap == null) {
          columnFamilyMap = new HashMap<>();
          batch.put(columnFamily, columnFamilyMap);
        }
  
        List<Byte> keyObject = byteArrayToList(key);
        // Copy the value bytes since they'll likely be modified in place.
        columnFamilyMap.put(keyObject, Arrays.copyOf(val, val.length));
      }
  
      @Override
      public void write() throws RocksDBException {
        /* Write everything accumulated in the batch in case the client depends on the writes no showing up in the DB
         * until batch.write is called.  Though is that even true when using a real RocksDB write batch? */
        for (Map.Entry<T, Map<List<Byte>, byte[]>> cfEntry : batch.entrySet()) {
          T cf = cfEntry.getKey();
          for (Map.Entry<List<Byte>, byte[]> kvEntry : cfEntry.getValue().entrySet()) {
            parent.put(cf, byteListToArray(kvEntry.getKey()), kvEntry.getValue());
          }
        }
      }
    }
  
    public static class MockRocksDBIterator extends RocksDBIterator {
      Map<List<Byte>, byte[]> index;
      List<List<Byte>> keys;
      int cursor = 0;
  
      public MockRocksDBIterator(Map<List<Byte>, byte[]> index) {
        super();
        this.index = index;
      }
  
      @Override
      public void reset() {
        List<List<Byte>> keys = new ArrayList<>(index.keySet());
        Collections.sort(keys, (a, b) -> {
          for (int i = 0; i < a.size() && i < b.size(); i++) {
            int cmp = a.get(i).compareTo(b.get(i));
            if (cmp != 0) {
              return cmp; // Return the comparison results for the first mismatch.
            }
          }
          return Integer.valueOf(a.size()).compareTo(b.size()); // If the values are all the same, prefer the shorter.
        });
        this.keys = keys;
        this.cursor = 0;
      }
  
      @Override
      public void seekToFirst() {
        reset();
      }
  
      @Override
      public void seekToLast() {
        this.cursor = this.keys.size() - 1;
      }
  
      @Override
      public void next() {
        this.cursor++;
      }
  
      @Override
      public void prev() {
        this.cursor--;
      }
  
      @Override
      public boolean isValid() {
        return this.cursor >= 0 && this.cursor < this.keys.size();
      }
  
      @Override
      public byte[] value() {
        return index.get(keys.get(cursor));
      }
  
      @Override
      public byte[] key() {
        return byteListToArray(this.keys.get(cursor));
      }
    }
  }
  