package cn.axzo.msg.center.utils;

import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @author yanglin
 */
@Slf4j
public class BatchController {

    private final String operation;
    private final AsyncRunTasks asyncTasks;
    private final Object request;
    private final AtomicReference<Exception> error = new AtomicReference<>();
    private final int parallelismThreshold;
    private int submitCount = 0;

    public BatchController(String operation, Object request,
                           ExecutorService executor) {
        this(operation, request, executor, 2);
    }

    public BatchController(String operation, Object request,
                           ExecutorService executor, int parallelismThreshold) {
        this.operation = operation;
        this.asyncTasks = new AsyncRunTasks(executor);
        this.request = request;
        this.parallelismThreshold = parallelismThreshold;
    }

    public void submit(Runnable task) {
        Runnable runner = () -> {
            Exception exception = error.get();
            if (exception != null) {
                warn("operation aborted because of other batch's error", exception);
                return;
            }
            try {
                task.run();
            } catch (Exception e) {
                setOperationFailed(e);
            }
        };
        submitCount++;
        if (submitCount <= parallelismThreshold) {
            info("using sync");
            runner.run();
            return;
        }
        info(String.format("using async because of submitCount > %d", parallelismThreshold));
        try {
            asyncTasks.runAsync(runner);
        } catch (Exception e) {
            setOperationFailed(e);
        }
    }

    private void setOperationFailed(Exception e) {
        warn("operation failed", e);
        error.compareAndSet(null, e);
    }

    public void awaitTermination(long timeout, TimeUnit unit) throws Exception {
        asyncTasks.awaitTermination(timeout, unit);
        Exception exception = error.get();
        if (exception != null) {
            warn("operation failed", exception);
            throw exception;
        }
        log.info("operation completed. operation={}", operation);
    }

    private void info(String message) {
        log.info("{}. operation={}, submitBatchSize={}, request={}",
                message, operation, submitCount, JSON.toJSONString(request));
    }

    private void warn(String message, Exception e) {
        log.info("{}. operation={}, submitBatchSize={}, request={}",
                message, operation, submitCount, JSON.toJSONString(request), e);
    }

}