package org.kaaproject.kaa.client.logging.memory;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

import org.kaaproject.kaa.client.logging.LogBlock;
import org.kaaproject.kaa.client.logging.LogRecord;
import org.kaaproject.kaa.client.logging.LogStorage;
import org.kaaproject.kaa.client.logging.LogStorageStatus;
import org.kaaproject.kaa.client.logging.memory.MemBucket.MemBucketState;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class MemLogStorage implements LogStorage, LogStorageStatus {

    private static final Logger LOG = LoggerFactory.getLogger(MemLogStorage.class);

    private static final long DEFAULT_MAX_STORAGE_SIZE = 16 * 1024 * 1024;
    private static final long DEFAULT_MAX_BUCKET_SIZE = 16 * 1024;
    private static final int DEFAULT_MAX_BUCKET_RECORD_COUNT = 256;

    private final long maxStorageSize;
    private long maxBucketSize;
    private int maxBucketRecordCount;

    private volatile long consumedVolume;
    private volatile long recordCount;

    private final AtomicInteger bucketIdSeq = new AtomicInteger();
    private MemBucket currentBucket;
    private final Map<Integer, MemBucket> buckets;
    
    public MemLogStorage() {
        this(DEFAULT_MAX_BUCKET_SIZE, DEFAULT_MAX_BUCKET_RECORD_COUNT);
    }
    
    public MemLogStorage(long bucketSize, int bucketRecordCount) {
        this(DEFAULT_MAX_STORAGE_SIZE, bucketSize, bucketRecordCount);
    }

    public MemLogStorage(long maxStorageSize, long bucketSize, int bucketRecordCount) {
        super();
        this.maxStorageSize = maxStorageSize;
        this.maxBucketSize = bucketSize;
        this.maxBucketRecordCount = bucketRecordCount;
        this.buckets = new HashMap<Integer, MemBucket>();
    }

    @Override
    public long getConsumedVolume() {
        LOG.debug("Consumed volume: {}", consumedVolume);
        return consumedVolume;
    }

    @Override
    public long getRecordCount() {
        LOG.debug("Record count: {}", recordCount);
        return recordCount;
    }

    @Override
    public void addLogRecord(LogRecord record) {
        LOG.trace("Adding new log record with size {}", record.getSize());
        if(record.getSize() > maxBucketSize){
            throw new IllegalArgumentException("Record size(" + record.getSize() + ") is bigger than max bucket size (" + maxBucketSize + ")!");
        }
        if (getConsumedVolume() + record.getSize() > maxStorageSize) {
            throw new IllegalStateException("Storage is full!");
        }
        synchronized (buckets) {
            if (currentBucket == null || currentBucket.getState() != MemBucketState.FREE) {
                currentBucket = new MemBucket(bucketIdSeq.getAndIncrement(), maxBucketSize, maxBucketRecordCount);
                buckets.put(currentBucket.getId(), currentBucket);
            }
            if (!currentBucket.addRecord(record)) {
                LOG.trace("Current bucket is full. Creating new one.");
                currentBucket.setState(MemBucketState.FULL);
                currentBucket = new MemBucket(bucketIdSeq.getAndIncrement(), maxBucketSize, maxBucketRecordCount);
                buckets.put(currentBucket.getId(), currentBucket);
                currentBucket.addRecord(record);
            }
            recordCount++;
            consumedVolume += record.getSize();
        }
        LOG.trace("Added a new log record to bucket [{}]", currentBucket.getId());
    }

    @Override
    public LogBlock getRecordBlock(long blockSize, int batchCount) {
        LOG.trace("Getting new record block with block size = {} and count = {}", blockSize, batchCount);
        if (blockSize > maxBucketSize || batchCount > maxBucketRecordCount) {
            //TODO: add support of block resize
            LOG.warn("Resize of record block is not supported yet");
        }

        LogBlock result = null;
        MemBucket bucketCandidate = null;
        synchronized (buckets) {
            for (MemBucket bucket : buckets.values()) {
                if (bucket.getState() == MemBucketState.FREE) {
                    bucketCandidate = bucket;
                }
                if (bucket.getState() == MemBucketState.FULL) {
                    bucket.setState(MemBucketState.PENDING);
                    bucketCandidate = bucket;
                    break;
                }
            }

            if (bucketCandidate != null) {
                consumedVolume -= bucketCandidate.getSize();
                recordCount -= bucketCandidate.getCount();
                if (bucketCandidate.getState() == MemBucketState.FREE) {
                    LOG.trace("Only a bucket with state FREE found: [{}]. Changing its state to PENDING", bucketCandidate.getId());
                    bucketCandidate.setState(MemBucketState.PENDING);
                }
                if (bucketCandidate.getSize() <= blockSize && bucketCandidate.getCount() <= batchCount) {
                    result = new LogBlock(bucketCandidate.getId(), bucketCandidate.getRecords());
                    LOG.debug("Return record block with records count: [{}]", bucketCandidate.getCount());
                } else {
                    LOG.debug("Shrinking bucket {} to new size: [{}] and count: [{}]", bucketCandidate, blockSize, batchCount);
                    List<LogRecord> overSized = bucketCandidate.shrinkToSize(blockSize, batchCount);
                    result = new LogBlock(bucketCandidate.getId(), bucketCandidate.getRecords());
                    for (LogRecord record : overSized) {
                        addLogRecord(record);
                    }
                }
            }
        }
        return result;
    }

    @Override
    public void removeRecordBlock(int id) {
        LOG.trace("Removing record block with id [{}]", id);
        synchronized (buckets) {
            if (buckets.remove(id) != null) {
                LOG.debug("Record block [{}] removed", id);
            } else {
                LOG.debug("Failed to remove record block [{}]", id);
            }
        }
    }

    @Override
    public void notifyUploadFailed(int id) {
        LOG.trace("Upload of record block [{}] failed", id);
        synchronized (buckets) {
            buckets.get(id).setState(MemBucketState.FULL);
            consumedVolume += buckets.get(id).getSize();
            recordCount += buckets.get(id).getCount();
        }
    }

    @Override
    public void close() {
        // automatically done by GC
    }

    @Override
    public LogStorageStatus getStatus() {
        return this;
    }

}

@@ -2266,18 +2266,19 @@
 ketSize)
+ 
 %7B%0A
-
         
@@ -2436,16 +2436,50 @@
     
-if (getC
+synchronized (buckets) %7B%0A            if (c
 onsu
@@ -2487,18 +2487,16 @@
 edVolume
-()
  + recor
@@ -2519,32 +2519,36 @@
 xStorageSize) %7B%0A
+    
             thro
@@ -2608,42 +2608,13 @@
     
-%7D%0A
     
-    synchronized (buckets) %7B
+%7D
 %0A   

fix conditional check to work concurrently