/*
 * Licensed to Elastic Search and Shay Banon under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. Elastic Search licenses this
 * file to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package org.elasticsearch.common.io.stream;

import org.elasticsearch.common.thread.ThreadLocals;

import java.io.IOException;
import java.util.Arrays;

/**
 * @author kimchy (shay.banon)
 */
public class BytesStreamOutput extends StreamOutput {

    /**
     * A thread local based cache of {@link BytesStreamOutput}.
     */
    public static class Cached {

        private static final ThreadLocal<ThreadLocals.CleanableValue<BytesStreamOutput>> cache = new ThreadLocal<ThreadLocals.CleanableValue<BytesStreamOutput>>() {
            @Override protected ThreadLocals.CleanableValue<BytesStreamOutput> initialValue() {
                return new ThreadLocals.CleanableValue<BytesStreamOutput>(new BytesStreamOutput());
            }
        };

        private static final ThreadLocal<ThreadLocals.CleanableValue<HandlesStreamOutput>> cacheHandles = new ThreadLocal<ThreadLocals.CleanableValue<HandlesStreamOutput>>() {
            @Override protected ThreadLocals.CleanableValue<HandlesStreamOutput> initialValue() {
                return new ThreadLocals.CleanableValue<HandlesStreamOutput>(new HandlesStreamOutput(new BytesStreamOutput()));
            }
        };

        /**
         * Returns the cached thread local byte stream, with its internal stream cleared.
         */
        public static BytesStreamOutput cached() {
            BytesStreamOutput os = cache.get().get();
            os.reset();
            return os;
        }

        public static HandlesStreamOutput cachedHandles() throws IOException {
            HandlesStreamOutput os = cacheHandles.get().get();
            os.reset();
            return os;
        }
    }

    /**
     * The buffer where data is stored.
     */
    protected byte buf[];

    /**
     * The number of valid bytes in the buffer.
     */
    protected int count;

    public BytesStreamOutput() {
        this(126);
    }

    public BytesStreamOutput(int size) {
        this.buf = new byte[size];
    }

    @Override public void writeByte(byte b) throws IOException {
        int newcount = count + 1;
        if (newcount > buf.length) {
            buf = Arrays.copyOf(buf, Math.max(buf.length << 1, newcount));
        }
        buf[count] = b;
        count = newcount;
    }

    @Override public void writeBytes(byte[] b, int offset, int length) throws IOException {
        if (length == 0) {
            return;
        }
        int newcount = count + length;
        if (newcount > buf.length) {
            buf = Arrays.copyOf(buf, Math.max(buf.length << 1, newcount));
        }
        System.arraycopy(b, offset, buf, count, length);
        count = newcount;
    }

    public void reset() {
        count = 0;
    }

    @Override public void flush() throws IOException {
        // nothing to do there
    }

    @Override public void close() throws IOException {
        // nothing to do here
    }

    /**
     * Creates a newly allocated byte array. Its size is the current
     * size of this output stream and the valid contents of the buffer
     * have been copied into it.
     *
     * @return the current contents of this output stream, as a byte array.
     * @see java.io.ByteArrayOutputStream#size()
     */
    public byte copiedByteArray()[] {
        return Arrays.copyOf(buf, count);
    }

    /**
     * Returns the underlying byte array. Note, use {@link #size()} in order to know
     * the length of it.
     */
    public byte[] unsafeByteArray() {
        return buf;
    }

    /**
     * Returns the current size of the buffer.
     *
     * @return the value of the <code>count</code> field, which is the number
     *         of valid bytes in this output stream.
     * @see java.io.ByteArrayOutputStream#count
     */
    public int size() {
        return count;
    }
}

@@ -1175,144 +1175,95 @@
     
-private static final ThreadLocal%3CThreadLocals.CleanableValue%3CBytesStreamOutput%3E%3E cache = new ThreadLocal%3CThreadLocals.CleanableValue%3CByt
+static class Entry %7B%0A            final BytesStreamOutput bytes;%0A            final Handl
 esSt
@@ -1264,38 +1264,42 @@
 dlesStreamOutput
-%3E%3E() %7B
+ handles;%0A
 %0A            @Ov
@@ -1299,190 +1299,144 @@
     
-@Override protected ThreadLocals.CleanableValue%3CBytesStreamOutput%3E initialValue() %7B%0A                return new ThreadLocals.CleanableValue%3CBytesStreamOutput%3E(new BytesStreamOutput())
+Entry(BytesStreamOutput bytes, HandlesStreamOutput handles) %7B%0A                this.bytes = bytes;%0A                this.handles = handles
 ;%0A  
@@ -1448,33 +1448,32 @@
      %7D%0A        %7D
-;
 %0A%0A        privat
@@ -1531,35 +1531,21 @@
 lue%3C
-HandlesStreamOutput
+Entry
 %3E%3E cache
 Hand
@@ -1540,23 +1540,16 @@
 %3E%3E cache
-Handles
  = new T
@@ -1587,35 +1587,21 @@
 leValue%3C
-HandlesStreamOutput
+Entry
 %3E%3E() %7B%0A 
@@ -1659,35 +1659,21 @@
 leValue%3C
-HandlesStreamOutput
+Entry
 %3E initia
@@ -1703,47 +1703,75 @@
     
-return new ThreadLocals.CleanableValue%3C
+BytesStreamOutput bytes = new BytesStreamOutput();%0A                
 Hand
@@ -1785,18 +1785,27 @@
 amOutput
-%3E(
+ handles = 
 new Hand
@@ -1820,39 +1820,110 @@
 mOutput(
-new BytesStreamOutput()
+bytes);%0A                return new ThreadLocals.CleanableValue%3CEntry%3E(new Entry(bytes, handles
 ));%0A    
@@ -2153,32 +2153,38 @@
 ache.get().get()
+.bytes
 ;%0A            os
@@ -2348,23 +2348,16 @@
  = cache
-Handles
 .get().g
@@ -2356,24 +2356,32 @@
 .get().get()
+.handles
 ;%0A          

improve thread cached bytes allocation