package com.cube.storage;

import java.util.*;
import java.util.concurrent.ConcurrentSkipListMap;

/**
 * In-memory sorted table for recent writes.
 */
public class MemTable {
    
    private final ConcurrentSkipListMap<String, byte[]> data;
    private volatile long size;
    
    public MemTable() {
        this.data = new ConcurrentSkipListMap<>();
        this.size = 0;
    }
    
    public void put(String key, byte[] value) {
        byte[] oldValue = data.put(key, value);
        
        if (oldValue != null) {
            size -= estimateEntrySize(key, oldValue);
        }
        size += estimateEntrySize(key, value);
    }
    
    public byte[] get(String key) {
        return data.get(key);
    }
    
    public boolean contains(String key) {
        return data.containsKey(key);
    }
    
    public Collection<String> scan(String prefix) {
        if (prefix == null || prefix.isEmpty()) {
            return new ArrayList<>(data.keySet());
        }
        
        List<String> result = new ArrayList<>();
        String startKey = prefix;
        String endKey = prefix + Character.MAX_VALUE;
        
        for (String key : data.subMap(startKey, endKey).keySet()) {
            result.add(key);
        }
        
        return result;
    }
    
    public Map<String, byte[]> scanEntries(String prefix) {
        if (prefix == null || prefix.isEmpty()) {
            return new TreeMap<>(data);
        }
        
        Map<String, byte[]> result = new TreeMap<>();
        String startKey = prefix;
        String endKey = prefix + Character.MAX_VALUE;
        
        for (Map.Entry<String, byte[]> entry : data.subMap(startKey, endKey).entrySet()) {
            result.put(entry.getKey(), entry.getValue());
        }
        
        return result;
    }
    
    public Map<String, byte[]> getEntries() {
        return new TreeMap<>(data);
    }
    
    public long size() {
        return size;
    }
    
    public int getKeyCount() {
        return data.size();
    }
    
    public boolean isEmpty() {
        return data.isEmpty();
    }
    
    public void clear() {
        data.clear();
        size = 0;
    }
    
    private long estimateEntrySize(String key, byte[] value) {
        return (key.length() * 2L) + value.length + 32;
    }
    
    @Override
    public String toString() {
        return "MemTable{keys=" + data.size() + ", size=" + size + " bytes}";
    }
}