/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch 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.compress;

import org.apache.lucene.util.LineFileDocs;
import org.apache.lucene.util.TestUtil;
import org.elasticsearch.common.io.stream.ByteBufferStreamInput;
import org.elasticsearch.common.io.stream.OutputStreamStreamOutput;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.settings.ImmutableSettings;
import org.elasticsearch.test.ElasticsearchTestCase;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.Random;
import java.util.concurrent.CountDownLatch;

/**
 * Test streaming compression (e.g. used for recovery)
 */
public class CompressedStreamTests extends ElasticsearchTestCase {
    
    @Override
    public void setUp() throws Exception {
        super.setUp();
        CompressorFactory.configure(ImmutableSettings.settingsBuilder().put("compress.default.type", "lzf").build());
    }

    public void testRandom() throws IOException {
        Random r = getRandom();
        for (int i = 0; i < 10; i++) {
            byte bytes[] = new byte[TestUtil.nextInt(r, 1, 400000)];
            r.nextBytes(bytes);
            doTest(bytes);
        }
    }
    
    public void testRandomThreads() throws Exception {
        final Random r = getRandom();
        int threadCount = TestUtil.nextInt(r, 2, 10);
        Thread[] threads = new Thread[threadCount];
        final CountDownLatch startingGun = new CountDownLatch(1);
        for (int tid=0; tid < threadCount; tid++) {
            final long seed = r.nextLong();
            threads[tid] = new Thread() {
                @Override
                public void run() {
                    try {
                        Random r = new Random(seed);
                        startingGun.await();
                        for (int i = 0; i < 10; i++) {
                            byte bytes[] = new byte[TestUtil.nextInt(r, 1, 100000)];
                            r.nextBytes(bytes);
                            doTest(bytes);
                        }
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }
                }
            };
            threads[tid].start();
        }
        startingGun.countDown();
        for (Thread t : threads) {
            t.join();
        }
    }
    
    public void testLineDocs() throws IOException {
        Random r = getRandom();
        LineFileDocs lineFileDocs = new LineFileDocs(r);
        for (int i = 0; i < 10; i++) {
            int numDocs = TestUtil.nextInt(r, 1, 200);
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            for (int j = 0; j < numDocs; j++) {
                String s = lineFileDocs.nextDoc().get("body");
                bos.write(s.getBytes(StandardCharsets.UTF_8));
            }
            doTest(bos.toByteArray());
        }
        lineFileDocs.close();
    }
    
    public void testLineDocsThreads() throws Exception {
        final Random r = getRandom();
        int threadCount = TestUtil.nextInt(r, 2, 10);
        Thread[] threads = new Thread[threadCount];
        final CountDownLatch startingGun = new CountDownLatch(1);
        for (int tid=0; tid < threadCount; tid++) {
            final long seed = r.nextLong();
            threads[tid] = new Thread() {
                @Override
                public void run() {
                    try {
                        Random r = new Random(seed);
                        startingGun.await();
                        LineFileDocs lineFileDocs = new LineFileDocs(r);
                        for (int i = 0; i < 10; i++) {
                            int numDocs = TestUtil.nextInt(r, 1, 200);
                            ByteArrayOutputStream bos = new ByteArrayOutputStream();
                            for (int j = 0; j < numDocs; j++) {
                                String s = lineFileDocs.nextDoc().get("body");
                                bos.write(s.getBytes(StandardCharsets.UTF_8));
                            }
                            doTest(bos.toByteArray());
                        }
                        lineFileDocs.close();
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }
                }
            };
            threads[tid].start();
        }
        startingGun.countDown();
        for (Thread t : threads) {
            t.join();
        }
    }
    
    public void testRepetitionsL() throws IOException {
        Random r = getRandom();
        for (int i = 0; i < 10; i++) {
            int numLongs = TestUtil.nextInt(r, 1, 10000);
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            long theValue = r.nextLong();
            for (int j = 0; j < numLongs; j++) {
                if (r.nextInt(10) == 0) {
                    theValue = r.nextLong();
                }
                bos.write((byte) (theValue >>> 56));
                bos.write((byte) (theValue >>> 48));
                bos.write((byte) (theValue >>> 40));
                bos.write((byte) (theValue >>> 32));
                bos.write((byte) (theValue >>> 24));
                bos.write((byte) (theValue >>> 16));
                bos.write((byte) (theValue >>> 8));
                bos.write((byte) theValue);
            }
            doTest(bos.toByteArray());
        }
    }
    
    public void testRepetitionsLThreads() throws Exception {
        final Random r = getRandom();
        int threadCount = TestUtil.nextInt(r, 2, 10);
        Thread[] threads = new Thread[threadCount];
        final CountDownLatch startingGun = new CountDownLatch(1);
        for (int tid=0; tid < threadCount; tid++) {
            final long seed = r.nextLong();
            threads[tid] = new Thread() {
                @Override
                public void run() {
                    try {
                        Random r = new Random(seed);
                        startingGun.await();
                        for (int i = 0; i < 10; i++) {
                            int numLongs = TestUtil.nextInt(r, 1, 10000);
                            ByteArrayOutputStream bos = new ByteArrayOutputStream();
                            long theValue = r.nextLong();
                            for (int j = 0; j < numLongs; j++) {
                                if (r.nextInt(10) == 0) {
                                    theValue = r.nextLong();
                                }
                                bos.write((byte) (theValue >>> 56));
                                bos.write((byte) (theValue >>> 48));
                                bos.write((byte) (theValue >>> 40));
                                bos.write((byte) (theValue >>> 32));
                                bos.write((byte) (theValue >>> 24));
                                bos.write((byte) (theValue >>> 16));
                                bos.write((byte) (theValue >>> 8));
                                bos.write((byte) theValue);
                            }
                            doTest(bos.toByteArray());
                        }
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }
                }
            };
            threads[tid].start();
        }
        startingGun.countDown();
        for (Thread t : threads) {
            t.join();
        }
    }
    
    public void testRepetitionsI() throws IOException {
        Random r = getRandom();
        for (int i = 0; i < 10; i++) {
            int numInts = TestUtil.nextInt(r, 1, 20000);
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            int theValue = r.nextInt();
            for (int j = 0; j < numInts; j++) {
                if (r.nextInt(10) == 0) {
                    theValue = r.nextInt();
                }
                bos.write((byte) (theValue >>> 24));
                bos.write((byte) (theValue >>> 16));
                bos.write((byte) (theValue >>> 8));
                bos.write((byte) theValue);
            }
            doTest(bos.toByteArray());
        }
    }
    
    public void testRepetitionsIThreads() throws Exception {
        final Random r = getRandom();
        int threadCount = TestUtil.nextInt(r, 2, 10);
        Thread[] threads = new Thread[threadCount];
        final CountDownLatch startingGun = new CountDownLatch(1);
        for (int tid=0; tid < threadCount; tid++) {
            final long seed = r.nextLong();
            threads[tid] = new Thread() {
                @Override
                public void run() {
                    try {
                        Random r = new Random(seed);
                        startingGun.await();
                        for (int i = 0; i < 10; i++) {
                            int numInts = TestUtil.nextInt(r, 1, 20000);
                            ByteArrayOutputStream bos = new ByteArrayOutputStream();
                            int theValue = r.nextInt();
                            for (int j = 0; j < numInts; j++) {
                                if (r.nextInt(10) == 0) {
                                    theValue = r.nextInt();
                                }
                                bos.write((byte) (theValue >>> 24));
                                bos.write((byte) (theValue >>> 16));
                                bos.write((byte) (theValue >>> 8));
                                bos.write((byte) theValue);
                            }
                            doTest(bos.toByteArray());
                        }
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }
                }
            };
            threads[tid].start();
        }
        startingGun.countDown();
        for (Thread t : threads) {
            t.join();
        }
    }
    
    public void testRepetitionsS() throws IOException {
        Random r = getRandom();
        for (int i = 0; i < 10; i++) {
            int numShorts = TestUtil.nextInt(r, 1, 40000);
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            short theValue = (short) r.nextInt(65535);
            for (int j = 0; j < numShorts; j++) {
                if (r.nextInt(10) == 0) {
                    theValue = (short) r.nextInt(65535);
                }
                bos.write((byte) (theValue >>> 8));
                bos.write((byte) theValue);
            }
            doTest(bos.toByteArray());
        }
    }

    public void testMixed() throws IOException {
        Random r = getRandom();
        LineFileDocs lineFileDocs = new LineFileDocs(r);
        for (int i = 0; i < 2; ++i) {
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            int prevInt = r.nextInt();
            long prevLong = r.nextLong();
            while (bos.size() < 400000) {
                switch (r.nextInt(4)) {
                case 0:
                    addInt(r, prevInt, bos);
                    break;
                case 1:
                    addLong(r, prevLong, bos);
                    break;
                case 2:
                    addString(lineFileDocs, bos);
                    break;
                case 3:
                    addBytes(r, bos);
                    break;
                default:
                    throw new IllegalStateException("Random is broken");
                }
            }
            doTest(bos.toByteArray());
        }
    }

    private void addLong(Random r, long prev, ByteArrayOutputStream bos) {
        long theValue = prev;
        if (r.nextInt(10) != 0) {
            theValue = r.nextLong();
        }
        bos.write((byte) (theValue >>> 56));
        bos.write((byte) (theValue >>> 48));
        bos.write((byte) (theValue >>> 40));
        bos.write((byte) (theValue >>> 32));
        bos.write((byte) (theValue >>> 24));
        bos.write((byte) (theValue >>> 16));
        bos.write((byte) (theValue >>> 8));
        bos.write((byte) theValue);
    }

    private void addInt(Random r, int prev, ByteArrayOutputStream bos) {
        int theValue = prev;
        if (r.nextInt(10) != 0) {
            theValue = r.nextInt();
        }
        bos.write((byte) (theValue >>> 24));
        bos.write((byte) (theValue >>> 16));
        bos.write((byte) (theValue >>> 8));
        bos.write((byte) theValue);
    }

    private void addString(LineFileDocs lineFileDocs, ByteArrayOutputStream bos) throws IOException {
        String s = lineFileDocs.nextDoc().get("body");
        bos.write(s.getBytes(StandardCharsets.UTF_8));
    }

    private void addBytes(Random r, ByteArrayOutputStream bos) throws IOException {
        byte bytes[] = new byte[TestUtil.nextInt(r, 1, 10000)];
        r.nextBytes(bytes);
        bos.write(bytes);
    }
    
    public void testRepetitionsSThreads() throws Exception {
        final Random r = getRandom();
        int threadCount = TestUtil.nextInt(r, 2, 10);
        Thread[] threads = new Thread[threadCount];
        final CountDownLatch startingGun = new CountDownLatch(1);
        for (int tid=0; tid < threadCount; tid++) {
            final long seed = r.nextLong();
            threads[tid] = new Thread() {
                @Override
                public void run() {
                    try {
                        Random r = new Random(seed);
                        startingGun.await();
                        for (int i = 0; i < 10; i++) {
                            int numShorts = TestUtil.nextInt(r, 1, 40000);
                            ByteArrayOutputStream bos = new ByteArrayOutputStream();
                            short theValue = (short) r.nextInt(65535);
                            for (int j = 0; j < numShorts; j++) {
                                if (r.nextInt(10) == 0) {
                                    theValue = (short) r.nextInt(65535);
                                }
                                bos.write((byte) (theValue >>> 8));
                                bos.write((byte) theValue);
                            }
                            doTest(bos.toByteArray());
                        }
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }
                }
            };
            threads[tid].start();
        }
        startingGun.countDown();
        for (Thread t : threads) {
            t.join();
        }
    }
       
    private void doTest(byte bytes[]) throws IOException {
        ByteBuffer bb = ByteBuffer.wrap(bytes);
        StreamInput rawIn = new ByteBufferStreamInput(bb);
        Compressor c = CompressorFactory.defaultCompressor();
        
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        OutputStreamStreamOutput rawOs = new OutputStreamStreamOutput(bos);
        StreamOutput os = c.streamOutput(rawOs);
        
        Random r = getRandom();
        int bufferSize = r.nextBoolean() ? 65535 : TestUtil.nextInt(getRandom(), 1, 70000);
        int prepadding = r.nextInt(70000);
        int postpadding = r.nextInt(70000);
        byte buffer[] = new byte[prepadding + bufferSize + postpadding];
        r.nextBytes(buffer); // fill block completely with junk
        int len;
        while ((len = rawIn.read(buffer, prepadding, bufferSize)) != -1) {
            os.write(buffer, prepadding, len);
        }
        os.close();
        rawIn.close();
        
        // now we have compressed byte array
        
        byte compressed[] = bos.toByteArray();
        ByteBuffer bb2 = ByteBuffer.wrap(compressed);
        StreamInput compressedIn = new ByteBufferStreamInput(bb2);
        StreamInput in = c.streamInput(compressedIn);
        
        // randomize constants again
        bufferSize = r.nextBoolean() ? 65535 : TestUtil.nextInt(getRandom(), 1, 70000);
        prepadding = r.nextInt(70000);
        postpadding = r.nextInt(70000);
        buffer = new byte[prepadding + bufferSize + postpadding];
        r.nextBytes(buffer); // fill block completely with junk
        
        ByteArrayOutputStream uncompressedOut = new ByteArrayOutputStream();
        while ((len = in.read(buffer, prepadding, bufferSize)) != -1) {
            uncompressedOut.write(buffer, prepadding, len);
        }
        uncompressedOut.close();
        
        assertArrayEquals(bytes, uncompressedOut.toByteArray());
    }
}
