package ru.yandex.webmaster3.worker.queue;

import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.datatype.joda.JodaModule;
import com.fasterxml.jackson.module.paramnames.ParameterNamesModule;
import org.apache.commons.io.FileUtils;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.jetbrains.annotations.NotNull;
import org.joda.time.DateTime;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import ru.yandex.webmaster3.core.data.WebmasterHostId;
import ru.yandex.webmaster3.core.http.WebmasterJsonModule;
import ru.yandex.webmaster3.core.util.IdUtils;
import ru.yandex.webmaster3.core.worker.task.WorkerTaskData;
import ru.yandex.webmaster3.core.worker.task.WorkerTaskPriority;
import ru.yandex.webmaster3.core.worker.task.WorkerTaskType;
import ru.yandex.webmaster3.worker.Task;
import ru.yandex.webmaster3.worker.TaskRegistry;

import java.io.*;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * @author aherman
 */
class TaskQueue {
    private static final Logger log = LoggerFactory.getLogger(TaskQueue.class);

    private static final ObjectMapper OM = new ObjectMapper()
            .disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
            .registerModules(new JodaModule(), new WebmasterJsonModule(false), new ParameterNamesModule());

    public static final String INFO_PREFIX = "taskQueue.";
    public static final String INFO_SUFFIX = ".info";

    public static final String DATA_PREFIX = "taskQueue.";
    public static final String DATA_SUFFIX = ".data";

    private AtomicInteger chunkIdCounter = new AtomicInteger();

    private int memoryQueueSize = 10;
    private int taskQueueChunkSize = 10000;

    private Lock writeLock = new ReentrantLock();
    private Lock readLock = new ReentrantLock();
    private Condition hasNewTasks = writeLock.newCondition();

    //task queue
    private Map<Pair<WorkerTaskType, WorkerTaskPriority>, QueueChunks> taskChunks = new HashMap<>();
    private final AtomicBoolean shutdown = new AtomicBoolean(false);

    private final TaskRegistry taskRegistry;
    private final File taskQueueFolder;
    private final TaskScheduler taskScheduler;
    private final TaskQueueMetrics taskQueueMetrics;

    @Autowired
    public TaskQueue(TaskRegistry taskRegistry,
                     @Value("${webmaster3.worker.taskQueue.folder}") File taskQueueFolder,
                     TaskScheduler taskScheduler,
                     TaskQueueMetrics taskQueueMetrics) {
        this.taskRegistry = taskRegistry;
        this.taskQueueFolder = taskQueueFolder;
        this.taskScheduler = taskScheduler;
        this.taskQueueMetrics = taskQueueMetrics;
    }

    public void init() throws IOException {
        if (!taskQueueFolder.exists()) {
            if (!taskQueueFolder.mkdirs()) {
                throw new IllegalStateException("Unable to create task queue folder: " + taskQueueFolder.getAbsolutePath());
            }
        } else if (!taskQueueFolder.isDirectory()) {
            throw new IllegalStateException(taskQueueFolder.getAbsolutePath() + " must be folder");
        }

        // Temporary
        cleanOldTaskFiles();

        for (WorkerTaskType type : taskRegistry.getTaskRegistryMap().keySet()) {
            int totalQueueSize = 0;
            for (WorkerTaskPriority priority : WorkerTaskPriority.values()) {
                QueueChunks queueChunks = new QueueChunks(type, priority, taskQueueFolder);
                totalQueueSize += queueChunks.getChunks().stream().mapToInt(QueueChunk::getChunkSize).sum();
                taskChunks.put(Pair.of(type, priority), queueChunks);
            }
            // общий размер очереди для любых приоритетов
            taskQueueMetrics.setQueueSize(type, totalQueueSize);
        }
    }

    private void cleanOldTaskFiles() {
        File[] files = taskQueueFolder.listFiles(File::isFile);
        for (File file : files) {
            file.delete();
        }
    }

    public void enqueueTask(TaskRunData taskRunData) throws IOException {
        if (shutdown.get()) {
            log.error("Queue shutdown, drop task: {}", taskRunData.getTaskId());
            return;
        }

        writeLock.lock();
        try {
            WorkerTaskType type = taskRunData.getTaskData().getTaskType();
            WorkerTaskPriority priority = taskRunData.getTaskPriority();
            QueueChunks chunks = taskChunks.get(Pair.of(type, priority));
            readLock.lock();
            try {
                QueueChunk writeChunk = chunks.getWriteChunk();

                if (writeChunk.addTask(taskRunData)) {
                    taskQueueMetrics.taskEnqueued(taskRunData.getTaskId().getTaskType());
                    hasNewTasks.signal();
                    return;
                }

                writeChunk = new QueueChunk(chunkIdCounter.getAndIncrement(), taskQueueChunkSize);
                writeChunk.addTask(taskRunData);
                taskQueueMetrics.taskEnqueued(taskRunData.getTaskId().getTaskType());

                if (chunks.getChunks().size() > memoryQueueSize) {
                    QueueChunk lastChunk = chunks.getChunks().peekLast();
                    chunks.unload(lastChunk);
                }
                chunks.setWriteChunk(writeChunk);
                hasNewTasks.signalAll();
            } finally {
                readLock.unlock();
            }
        } finally {
            writeLock.unlock();
        }
    }

    public TaskRunData pollTask() throws InterruptedException {
        if (shutdown.get()) {
            return null;
        }
        return pollTask(true);
    }

    TaskRunData pollNoWait() throws InterruptedException {
        return pollTask(false);
    }

    private TaskRunData pollTask(boolean wait) throws InterruptedException {
        while (true) {
            readLock.lock();
            try {
                WorkerTaskType type = taskScheduler.pollTaskType();
                if (type != null) {
                    // сперва обрабатываем задачи с высоким приоритетом
                    for (WorkerTaskPriority priority : WorkerTaskPriority.values()) {
                        QueueChunks chunks = taskChunks.get(Pair.of(type, priority));

                        QueueChunk chunkInfo;
                        while ((chunkInfo = chunks.getChunks().peekFirst()) != null) {
                            QueueChunk.ChunkState state = chunkInfo.getState();
                            if (state == QueueChunk.ChunkState.UNLOADED) {
                                chunks.loadChunk(chunkInfo);
                            }

                            if (chunkInfo.hasNextTask()) {
                                TaskRunData result = chunkInfo.pollTask();
                                taskQueueMetrics.taskPolled(result.getTaskId());
                                return result;
                            }

                            if (chunkInfo.isFull()) {
                                chunks.getChunks().removeFirst();
                                chunks.cleanChunk(chunkInfo);
                                continue;
                            }

                            break;
                        }
                    }
                }
            } finally {
                readLock.unlock();
            }

            if (wait) {
                writeLock.lock();
                try {
                    hasNewTasks.await(500, TimeUnit.MILLISECONDS);
                } finally {
                    writeLock.unlock();
                }
            } else {
                return null;
            }
        }
    }

    void clear() {
        writeLock.lock();
        try {
            readLock.lock();
            try {
                for (Map.Entry<Pair<WorkerTaskType, WorkerTaskPriority>, QueueChunks> e : taskChunks.entrySet()) {
                    e.getValue().clear();
                    taskQueueMetrics.queueCleared(e.getKey().getLeft());
                }
            } finally {
                readLock.unlock();
            }
        } finally {
            writeLock.unlock();
        }
    }

    void clear(WorkerTaskType type) {
        writeLock.lock();
        try {
            readLock.lock();
            try {
                for (WorkerTaskPriority priority : WorkerTaskPriority.values()) {
                    taskChunks.get(Pair.of(type, priority)).clear();
                }
                taskQueueMetrics.queueCleared(type);
            } finally {
                readLock.unlock();
            }
        } finally {
            writeLock.unlock();
        }
    }

    public void beginShutdown() {
        shutdown.set(true);
    }

    public void shutdown() throws IOException {
        log.info("Save unfinished tasks to: {}", taskQueueFolder.getAbsolutePath());
        int taskCount = 0;

        try {
            for (Map.Entry<Pair<WorkerTaskType, WorkerTaskPriority>, QueueChunks> e : taskChunks.entrySet()) {
                for (QueueChunk taskChunk : e.getValue().getChunks()) {
                    taskCount += taskChunk.getChunkSize();
                    e.getValue().unload(taskChunk);
                }
            }
        } catch (Exception e) {
            log.error("Unable to save tasks", e);
        }
        log.info("{}/{} unfinished tasks saved", taskCount, taskChunks.size());
    }

    /**
     * @author tsyplyaev
     */
    class QueueChunks {
        WorkerTaskType type;
        WorkerTaskPriority priority;
        File folder;
        ArrayDeque<QueueChunk> chunks = new ArrayDeque<>();
        QueueChunk writeChunk;

        public QueueChunks(WorkerTaskType type, WorkerTaskPriority priority, File baseFolder) throws IOException {
            this.type = type;
            this.priority = priority;

            this.folder = new File(baseFolder, type.toString() + priority.getPathSuffix());
            if (!folder.exists()) {
                if (!folder.mkdirs()) {
                    throw new IllegalStateException("Unable to create task queue folder: " + folder.getAbsolutePath());
                }
            } else if (!folder.isDirectory()) {
                throw new IllegalStateException(folder.getAbsolutePath() + " must be folder");
            }

            log.info("Load unfinished tasks from: {}", folder.getAbsolutePath());
            this.chunks = loadChunkList(folder);

            log.info("Unfinished tasks loaded");

            this.writeChunk = new QueueChunk(chunkIdCounter.getAndIncrement(), taskQueueChunkSize);
            this.chunks.addLast(writeChunk);
        }

        public WorkerTaskType getType() {
           return type;
        }

        public WorkerTaskPriority getPriority() {
            return priority;
        }

        public File getFolder() {
            return folder;
        }

        public ArrayDeque<QueueChunk> getChunks() {
            return chunks;
        }

        public QueueChunk getWriteChunk() {
            return writeChunk;
        }

        public void setWriteChunk(QueueChunk writeChunk) {
            this.writeChunk = writeChunk;
            this.chunks.addLast(writeChunk);
        }

        public void clear() {
            this.writeChunk = new QueueChunk(chunkIdCounter.getAndIncrement(), taskQueueChunkSize);
            this.chunks.clear();
            this.chunks.add(this.writeChunk);

            try {
                FileUtils.cleanDirectory(folder);
            } catch (Exception e) {
                log.error("Unable to clear queue {}/{}", type, priority);
            }
        }

        //write
        @NotNull
        private File getChunkDataFile(QueueChunk lastChunk) {
            return new File(folder, DATA_PREFIX + lastChunk.getChunkId() + DATA_SUFFIX);
        }

        @NotNull
        private File getChunkInfoFile(QueueChunk lastChunk) {
            return new File(folder, INFO_PREFIX + lastChunk.getChunkId() + INFO_SUFFIX);
        }

        private void unload(QueueChunk chunkInfo) throws IOException {
            if (chunkInfo.isEmpty()) {
                // Skip empty chunks
                chunkInfo.setState(QueueChunk.ChunkState.UNLOADED, null);
                return;
            }

            File chunkInfoFile = getChunkInfoFile(chunkInfo);
            QueueChunk.ChunkState state = chunkInfo.getState();

            if (state == QueueChunk.ChunkState.WRITABLE) {
                File chunkDataFile = getChunkDataFile(chunkInfo);
                save(chunkInfo, chunkDataFile);
            }
            chunkInfo.setState(QueueChunk.ChunkState.UNLOADED, null);
            String infoLine = String.format("[%s] %s", DateTime.now(), QueueChunk.BitSetUtils.toBase64String(chunkInfo.getFlags()));
            switch (state) {
                case WRITABLE:
                    try (Writer w = new FileWriter(chunkInfoFile, false)) {
                        w.write(infoLine);
                        w.write('\n');
                    }
                    break;

                case LOADED:
                    try (Writer w = new FileWriter(chunkInfoFile, true)) {
                        w.write(infoLine);
                        w.write('\n');
                    }
                    break;
            }
        }

        void save(QueueChunk chunkInfo, File chunkDataFile) throws IOException {
            TaskRunData[] taskList = chunkInfo.getTaskList();
            try (BufferedWriter w = new BufferedWriter(new FileWriter(chunkDataFile, false))) {
                w.write(Integer.toString(chunkInfo.getTail()));
                w.write('\n');

                for (int i = 0; i < chunkInfo.getTail(); i++) {
                    TaskRunData taskRunData = taskList[i];
                    TaskId taskId = taskRunData.getTaskId();
                    String s = ""
                            + taskId.getTaskType().name()
                            + "\t"
                            + (taskId.getHostId() != null ? taskId.getHostId().toStringId() : "-")
                            + "\t"
                            + taskId.getTaskUUID()
                            + "\t"
                            + taskRunData.getExecutorType().name()
                            + "\t"
                            + OM.writeValueAsString(taskRunData.getTaskData());

                    w.write(Integer.toString(s.length()));
                    w.write('\n');
                    w.write(s);
                    w.write('\n');
                }
            }
        }

        private ArrayDeque<QueueChunk> loadChunkList(File folder) throws IOException {
            File[] files = folder.listFiles(f ->
                    f.isFile()
                            && f.getName().startsWith(INFO_PREFIX)
                            && f.getName().endsWith(INFO_SUFFIX)
            );
            if (files != null && files.length > 0) {
                ArrayList<QueueChunk> result = new ArrayList<>(files.length);
                for (File file : files) {
                    QueueChunk chunkInfo = loadChunkInfo(file);
                    if (chunkInfo.isEmpty()) {
                        cleanChunk(chunkInfo);
                    } else {
                        result.add(chunkInfo);
                    }
                }
                result.sort(
                        Comparator.comparingLong(QueueChunk::getChunkDateMs)
                                .thenComparing(
                                        Comparator.comparingInt(QueueChunk::getChunkCount)
                                )
                );
                return new ArrayDeque<>(result);
            } else {
                return new ArrayDeque<>(memoryQueueSize);
            }
        }

        QueueChunk loadChunkInfo(File file) throws IOException {
            String name = file.getName();
            String idStr = name.substring(INFO_PREFIX.length(), name.length() - INFO_SUFFIX.length());
            String[] id = StringUtils.split(idStr, '.');
            long chunkDateMs = Long.parseLong(id[0]);
            int chunkCount = Integer.parseInt(id[1]);
            String lastLine;
            try (BufferedReader br = new BufferedReader(new FileReader(file))) {
                List<String> lines = IOUtils.readLines(br);
                if (lines.isEmpty()) {
                    log.error("Empty chunk info, skip: {}", file.getAbsolutePath());
                    return new QueueChunk(chunkDateMs, chunkCount, new BitSet(0));
                }
                lastLine = lines.get(lines.size() - 1);
            }
            String[] split = StringUtils.split(lastLine, ' ');
            QueueChunk chunkInfo;
            if (split.length == 3) {
                chunkInfo = new QueueChunk(chunkDateMs, chunkCount, Integer.parseInt(split[1]), Integer.parseInt(split[2]));
            } else if (split.length == 2) {
                chunkInfo = new QueueChunk(chunkDateMs, chunkCount, QueueChunk.BitSetUtils.fromBase64String(split[1]));
            } else {
                chunkInfo = new QueueChunk(chunkDateMs, chunkCount, new BitSet(0));
            }

            chunkInfo.setState(QueueChunk.ChunkState.UNLOADED, null);
            return chunkInfo;
        }

        private void loadChunk(QueueChunk chunkInfo) {
            File chunkDataFile = getChunkDataFile(chunkInfo);
            TaskRunData[] taskRunDatas;
            try {
                taskRunDatas = loadChunk(chunkDataFile);
                if (taskRunDatas == null) {
                    return;
                }

                for (int i = 0; i < taskRunDatas.length; i++) {
                    TaskRunData taskRunData = taskRunDatas[i];
                    if (!StringUtils.isEmpty(taskRunData.getTaskDataStr())) {
                        Task task = taskRegistry.getTaskRegistryMap().get(taskRunData.getTaskId().getTaskType());
                        if (task == null) {
                            log.error("Unknown task type: {}", taskRunData.getTaskId());
                            taskRunDatas[i] = null;
                        } else {
                            WorkerTaskData taskData =
                                    (WorkerTaskData) OM.readValue(taskRunData.getTaskDataStr(), task.getDataClass());
                            taskRunDatas[i] = new TaskRunData(taskRunData.getTaskId(), taskRunData.getExecutorType(),
                                    taskData, priority);
                        }
                    }
                }
            } catch (IOException e) {
                log.error("Unable to load task data: {}", chunkDataFile.getAbsolutePath(), e);
                return;
            }

            chunkInfo.setState(QueueChunk.ChunkState.LOADED, taskRunDatas);
        }

        TaskRunData[] loadChunk(File chunkDataFile) throws IOException {
            TaskRunData[] result;
            try (BufferedReader br = new BufferedReader(new FileReader(chunkDataFile))) {
                int totalTaskCount = Integer.parseInt(br.readLine());
                if (totalTaskCount <= 0 || totalTaskCount > 1000000) {
                    log.error("Broken task data file: {}", chunkDataFile.getAbsolutePath());
                    return null;
                }

                result = new TaskRunData[totalTaskCount];

                for (int i = 0; i < totalTaskCount; i++) {
                    String sizeStr = br.readLine();
                    int size = Integer.parseInt(sizeStr);
                    char[] buffer = new char[size];
                    int actualSize = br.read(buffer);
                    if (actualSize != size) {
                        log.error("Broken chunk data: {}", chunkDataFile.getAbsolutePath());
                        return null;
                    }
                    char lastCh = (char) br.read();
                    if ('\n' != lastCh) {
                        log.error("Broken chunk data: {}", chunkDataFile.getAbsolutePath());
                        return null;
                    }
                    String s = new String(buffer);
                    String[] parts = StringUtils.split(s, "\t", 5);
                    TaskId taskId = createTaskId(parts[0], parts[1], parts[2]);
                    result[i] = new TaskRunData(taskId, TaskRunType.R.valueOfOrUnknown(parts[3]), parts[4]);
                }
            }
            return result;
        }

        TaskId createTaskId(String taskTypeStr, String hostIdStr, String taskUUIDStr) {
            WorkerTaskType taskType = WorkerTaskType.R.valueOfOrUnknown(taskTypeStr);
            WebmasterHostId hostId = "-".equals(hostIdStr) ? null : IdUtils.stringToHostId(hostIdStr);
            return new TaskId(taskType, hostId, UUID.fromString(taskUUIDStr));
        }

        void cleanChunk(QueueChunk chunkInfo) {
            log.info("Delete chunk: {}", chunkInfo.getChunkId());
            File chunkDataFile = getChunkDataFile(chunkInfo);
            File chunkInfoFile = getChunkInfoFile(chunkInfo);
            chunkDataFile.delete();
            chunkInfoFile.delete();
        }
    }

    public void setMemoryQueueSize(int memoryQueueSize) {
        this.memoryQueueSize = memoryQueueSize;
    }

    public void setTaskQueueChunkSize(int taskQueueChunkSize) {
        this.taskQueueChunkSize = taskQueueChunkSize;
    }
}
