package ru.yandex.webmaster3.core.worker.client;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.EnumMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import javax.annotation.PostConstruct;

import com.google.common.collect.Lists;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import ru.yandex.webmaster3.core.solomon.metric.SolomonCounter;
import ru.yandex.webmaster3.core.solomon.metric.SolomonKey;
import ru.yandex.webmaster3.core.solomon.metric.SolomonMetricConfiguration;
import ru.yandex.webmaster3.core.solomon.metric.SolomonMetricRegistry;
import ru.yandex.webmaster3.core.util.json.JsonMapping;
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.core.worker.task.model.WorkerTaskDataBatch;
import ru.yandex.webmaster3.core.worker.task.model.WorkerTaskDataWrapper;

/**
 * @author: ishalaru
 * DATE: 06.08.2019
 */
@Slf4j
@Component("lbWorkerClient")
public class LogBrokerWorkerClient implements WorkerClient {
    private static final String SOLOMON_ACTION_LABEL = "task";
    private static final String SOLOMON_ACTION_TYPE = "task_send_in_lb";
    private static final String SOLOMON_LABEL_DATA_TYPE = "task_type";
    private static final String SOLOMON_LABEL_COUNT = "count";
    private static final String SOLOMON_LABEL_DATA_TYPE_ACTION = "send";
    private static final int TASK_BATCH_SIZE = 200;
    private final BlockingQueue<WorkerTaskData> sendQueue = new ArrayBlockingQueue<>(5000);
    private final LogbrokerMultiTopicClient logbrokerMultiTopicClient;
    private final SolomonMetricRegistry solomonMetricRegistry;
    @Value("${webmaster3.core.lb.worker.client.solomon.enabled:false}")
    private boolean metricIsEnable;
    private SolomonMetricConfiguration solomonMetricConfiguration;
    private EnumMap<WorkerTaskType, SolomonCounter> solomonMetricMap;
    private ExecutorService sendQueueExecutor;

    @Autowired
    public LogBrokerWorkerClient(@Qualifier("logBrokerMultiTopicClient") LogbrokerMultiTopicClient logBrokerMultiTopicClient,
                                 SolomonMetricRegistry solomonMetricRegistry) {
        this.logbrokerMultiTopicClient = logBrokerMultiTopicClient;
        this.solomonMetricRegistry = solomonMetricRegistry;
    }

    @PostConstruct
    public void init() {
        solomonMetricConfiguration = new SolomonMetricConfiguration();
        solomonMetricConfiguration.setEnable(metricIsEnable);
        solomonMetricMap = new EnumMap<>(WorkerTaskType.class);
        for (WorkerTaskType value : WorkerTaskType.values()) {
            solomonMetricMap.put(value, createSolomonCounter(value.toString()));
        }
        sendQueueExecutor = Executors.newSingleThreadExecutor(
                new ThreadFactoryBuilder()
                        .setNameFormat("worker-client-queue-%d")
                        .setDaemon(true)
                        .build()
        );
        sendQueueExecutor.execute(new LogBrokerWorkerClient.SendWorker());
        log.info("LogBrokerWorkerClient init");
    }

    private SolomonCounter createSolomonCounter(String taskType) {
        SolomonKey baseKey = SolomonKey.create(SOLOMON_ACTION_LABEL, SOLOMON_ACTION_TYPE).withLabel(SOLOMON_LABEL_DATA_TYPE, taskType);
        return solomonMetricRegistry.createCounter(solomonMetricConfiguration, baseKey.withLabel(SOLOMON_LABEL_DATA_TYPE_ACTION, SOLOMON_LABEL_COUNT));
    }

    @Override
    public <TD extends WorkerTaskData> void enqueueTask(TD taskData) {
        UUID taskId = taskData.getTaskId();
        WorkerTaskType taskType = taskData.getTaskType();
        WorkerTaskPriority taskPriority = taskData.getTaskPriority();

        log.trace("Send task: id={} type={} {} queueSize={} taskPriority={}", taskId, taskType, taskData.getShortDescription(),
                sendQueue.size(), taskPriority);
        boolean enqueued;
        try {
            enqueued = sendQueue.offer(taskData, 1, TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            enqueued = false;
        }
        if (!enqueued) {
            log.error("Unable to add task: id={} type={} {}", taskId, taskType, taskData.getShortDescription());
        }
    }

    @Override
    public <TD extends WorkerTaskData> void enqueueBatch(Collection<TD> batch) {
        internalEnqueue(batch);
    }


    @Override
    public <TD extends WorkerTaskData> boolean checkedEnqueueTask(TD taskData) {
        UUID taskId = taskData.getTaskId();
        WorkerTaskType taskType = taskData.getTaskType();
        WorkerTaskPriority taskPriority = taskData.getTaskPriority();

        log.trace("Send task: id={} type={} {} queueSize={} taskPriority={}", taskId, taskType, taskData.getShortDescription(),
                sendQueue.size(), taskPriority);
        return internalEnqueue(Collections.singletonList(taskData));
    }

    @Override
    public <TD extends WorkerTaskData> boolean checkedEnqueueBatch(Collection<TD> batch) {
        return internalEnqueue(batch);
    }


    private boolean internalEnqueue(Collection<? extends WorkerTaskData> taskDataList) {
        log.trace("Send {} tasks: {}", taskDataList.size(), taskDataList);
        if (taskDataList.isEmpty()) {
            return true;
        }
        if (taskDataList.iterator().next().shouldRunLocally()) {
            log.warn("This type of protocol don't support locally run mode");
        }
        taskDataList.forEach(e -> solomonMetricMap.get(e.getTaskType()).update());
        final Map<WorkerTaskPriority, List<WorkerTaskDataWrapper>> collect = taskDataList.stream().
                map(WorkerTaskDataWrapper::fromWorkerTaskData)
                .collect(Collectors.groupingBy(e -> e.getTaskPriority(), Collectors.mapping(e -> e, Collectors.toList())));
        try {
            for (Map.Entry<WorkerTaskPriority, List<WorkerTaskDataWrapper>> item : collect.entrySet()) {
                if (item.getValue().size() > TASK_BATCH_SIZE) {
                    //Чисто защита от слишком толстых пачек
                    final List<List<WorkerTaskDataWrapper>> partition = Lists.partition(item.getValue(), TASK_BATCH_SIZE);
                    for (List<WorkerTaskDataWrapper> workerTaskDataWrappers : partition) {
                        final UUID uuid = UUID.randomUUID();
                        String array = JsonMapping.writeValueAsString(new WorkerTaskDataBatch(workerTaskDataWrappers, uuid));
                        logbrokerMultiTopicClient.write(array, item.getKey());

                    }
                } else {
                    final UUID uuid = UUID.randomUUID();
                    String array = JsonMapping.writeValueAsString(new WorkerTaskDataBatch(item.getValue(), uuid));
                    logbrokerMultiTopicClient.write(array, item.getKey());
                }
            }
        } catch (Exception exp) {
            log.error("Can't send message to workers.", exp);
            return false;
        }

        return true;
    }

    private class SendWorker implements Runnable {
        @Override
        public void run() {
            log.info("Start task");
            while (!Thread.interrupted()) {
                WorkerTaskData taskData;
                try {
                    taskData = sendQueue.take();
                } catch (InterruptedException e) {
                    break;
                }
                List<WorkerTaskData> taskDataList = new ArrayList<>(TASK_BATCH_SIZE);
                taskDataList.add(taskData);
                // затягиваем до 99 элементов (1 уже есть)
                sendQueue.drainTo(taskDataList, TASK_BATCH_SIZE - 1);
                internalEnqueue(taskDataList);

            }
            log.info("Stop task");
        }
    }


}
