/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF 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.apache.kafka.jmh.log;

import org.apache.kafka.clients.consumer.OffsetOutOfRangeException;
import org.apache.kafka.common.compress.Compression;
import org.apache.kafka.common.config.TopicConfig;
import org.apache.kafka.common.record.internal.FileLogInputStream;
import org.apache.kafka.common.record.internal.FileRecords;
import org.apache.kafka.common.record.internal.MemoryRecords;
import org.apache.kafka.common.record.internal.RecordBatch;
import org.apache.kafka.common.record.internal.Records;
import org.apache.kafka.common.record.internal.SimpleRecord;
import org.apache.kafka.common.utils.Exit;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.server.common.RequestLocal;
import org.apache.kafka.server.storage.log.FetchIsolation;
import org.apache.kafka.server.util.MockTime;
import org.apache.kafka.storage.internals.log.AppendOrigin;
import org.apache.kafka.storage.internals.log.FetchDataInfo;
import org.apache.kafka.storage.internals.log.LogAppendInfo;
import org.apache.kafka.storage.internals.log.LogConfig;
import org.apache.kafka.storage.internals.log.LogDirFailureChannel;
import org.apache.kafka.storage.internals.log.LogOffsetsListener;
import org.apache.kafka.storage.internals.log.ProducerStateManagerConfig;
import org.apache.kafka.storage.internals.log.UnifiedLog;
import org.apache.kafka.storage.internals.log.VerificationGuard;
import org.apache.kafka.storage.log.metrics.BrokerTopicStats;
import org.apache.kafka.test.TestUtils;

import java.io.File;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.Optional;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

import static org.apache.kafka.server.common.TransactionVersion.TV_UNKNOWN;


public class StressTestLog {
    private static final AtomicBoolean RUNNING = new AtomicBoolean(true);

    public static void main(String[] args) throws Exception {
        File tmp = TestUtils.tempDirectory();
        File dir = TestUtils.randomPartitionLogDir(tmp);
        MockTime time = new MockTime();
        Properties logProperties = new Properties();
        logProperties.put(TopicConfig.SEGMENT_BYTES_CONFIG, 64 * 1024 * 1024);
        logProperties.put(TopicConfig.MAX_MESSAGE_BYTES_CONFIG, Integer.MAX_VALUE);
        logProperties.put(TopicConfig.SEGMENT_INDEX_BYTES_CONFIG, 1024 * 1024);

        int fiveMinutesInMillis = (int) Duration.ofMinutes(5).toMillis();
        UnifiedLog log = UnifiedLog.create(
            dir,
            new LogConfig(logProperties),
            0L,
            0L,
            time.scheduler,
            new BrokerTopicStats(),
            time,
            fiveMinutesInMillis,
            new ProducerStateManagerConfig(600000, false),
            fiveMinutesInMillis,
            new LogDirFailureChannel(10),
            true,
            Optional.empty(),
            new ConcurrentHashMap<>(),
            false,
            LogOffsetsListener.NO_OP_OFFSETS_LISTENER
        );

        WriterThread writer = new WriterThread(log);
        writer.start();
        ReaderThread reader = new ReaderThread(log);
        reader.start();

        Exit.addShutdownHook("stress-test-shutdown-hook", () -> {
            try {
                RUNNING.set(false);
                writer.join();
                reader.join();
                Utils.delete(dir);
            } catch (Exception e) {
                e.printStackTrace();
            }
        });

        while (RUNNING.get()) {
            Thread.sleep(1000);
            System.out.printf("Reader offset = %d, writer offset = %d%n", reader.currentOffset, writer.currentOffset);
            writer.checkProgress();
            reader.checkProgress();
        }
    }

    abstract static class WorkerThread extends Thread {
        protected long currentOffset = 0;
        private long lastOffsetCheckpointed = currentOffset;
        private long lastProgressCheckTime = System.currentTimeMillis();

        @Override
        public void run() {
            try {
                while (StressTestLog.RUNNING.get()) {
                    work();
                }
            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                StressTestLog.RUNNING.set(false);
            }
        }

        protected abstract void work() throws Exception;

        public boolean isMakingProgress() {
            if (currentOffset > lastOffsetCheckpointed) {
                lastOffsetCheckpointed = currentOffset;
                return true;
            }
            return false;
        }

        public void checkProgress() {
            long curTime = System.currentTimeMillis();
            if ((curTime - lastProgressCheckTime) > 500) {
                if (!isMakingProgress()) {
                    throw new RuntimeException("Thread not making progress");
                }
                lastProgressCheckTime = curTime;
            }
        }
    }

    static class WriterThread extends WorkerThread {
        private final UnifiedLog log;

        public WriterThread(UnifiedLog log) {
            this.log = log;
        }

        @Override
        protected void work() throws Exception {
            byte[] value = Long.toString(currentOffset).getBytes(StandardCharsets.UTF_8);
            MemoryRecords records = MemoryRecords.withRecords(
                    RecordBatch.CURRENT_MAGIC_VALUE,
                    Compression.NONE,
                    new SimpleRecord(RecordBatch.NO_TIMESTAMP, null, value));
            LogAppendInfo logAppendInfo = log.appendAsLeader(records,
                    0,
                    AppendOrigin.CLIENT,
                    RequestLocal.noCaching(),
                    VerificationGuard.SENTINEL,
                    TV_UNKNOWN);

            if ((logAppendInfo.firstOffset() != -1 && logAppendInfo.firstOffset() != currentOffset)
                || logAppendInfo.lastOffset() != currentOffset) {
                throw new RuntimeException("Offsets do not match");
            }

            currentOffset++;
            if (currentOffset % 1000 == 0) {
                Thread.sleep(50);
            }
        }
    }

    static class ReaderThread extends WorkerThread {
        private final UnifiedLog log;

        public ReaderThread(UnifiedLog log) {
            this.log = log;
        }

        @Override
        protected void work() throws Exception {
            try {
                FetchDataInfo fetchDataInfo = log.read(
                    currentOffset,
                    1,
                    FetchIsolation.LOG_END,
                    true
                );

                Records records = fetchDataInfo.records;

                if (records instanceof FileRecords && records.sizeInBytes() > 0) {
                    FileLogInputStream.FileChannelRecordBatch first = ((FileRecords) records).batches().iterator().next();
                    if (first.lastOffset() != currentOffset) {
                        throw new RuntimeException("Expected to read the message at offset " + currentOffset);
                    }
                    if (first.sizeInBytes() != records.sizeInBytes()) {
                        throw new RuntimeException(String.format("Expected %d but got %d.", first.sizeInBytes(), records.sizeInBytes()));
                    }
                    currentOffset++;
                }
            } catch (OffsetOutOfRangeException e) {
                // this is ok
            }
        }
    }
}
