package com.alibaba.datax.plugin.reader.otsstreamreader.internal;

import com.alibaba.datax.common.exception.DataXException;
import com.alibaba.datax.common.plugin.RecordSender;
import com.alibaba.datax.plugin.reader.otsstreamreader.internal.config.OTSStreamReaderConfig;
import com.alibaba.datax.plugin.reader.otsstreamreader.internal.config.OTSStreamReaderConstants;
import com.alibaba.datax.plugin.reader.otsstreamreader.internal.core.*;
import com.alibaba.datax.plugin.reader.otsstreamreader.internal.model.ShardCheckpoint;
import com.alibaba.datax.plugin.reader.otsstreamreader.internal.model.StreamJob;
import com.alibaba.datax.plugin.reader.otsstreamreader.internal.utils.OTSHelper;
import com.alibaba.datax.plugin.reader.otsstreamreader.internal.utils.TimeUtils;
import com.alicloud.openservices.tablestore.*;
import com.alicloud.openservices.tablestore.model.*;
import com.aliyun.openservices.ots.internal.streamclient.model.CheckpointPosition;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

public class OTSStreamReaderSlaveProxy {
    private static final Logger LOG = LoggerFactory.getLogger(OTSStreamReaderSlaveProxy.class);
    private static AtomicInteger slaveNumber = new AtomicInteger(0);

    private OTSStreamReaderConfig config;
    private SyncClientInterface ots;
    private Map<String, ShardCheckpoint> shardToCheckpointMap = new ConcurrentHashMap<String, ShardCheckpoint>();
    private CheckpointTimeTracker checkpointInfoTracker;
    private OTSStreamReaderChecker checker;

    private StreamJob streamJob;
    private Map<String, StreamShard> allShardsMap; // all shards from job master
    private Map<String, StreamShard> ownedShards; // shards to read arranged by job master
    private boolean findCheckpoints; // whether find checkpoint for last job, if so, we should read from checkpoint and skip nothing.
    private String slaveId = UUID.randomUUID().toString();
    private StreamDetails streamDetails;

    public void init(final OTSStreamReaderConfig otsStreamReaderConfig, StreamJob streamJob, List<StreamShard> allShards, Set<String> ownedShardIds) {
        slaveNumber.getAndIncrement();
        this.config = otsStreamReaderConfig;
        this.ots = OTSHelper.getOTSInstance(config);
        this.streamJob = streamJob;
        this.streamDetails = OTSHelper.getStreamDetails(ots, this.streamJob.getTableName());
        this.checkpointInfoTracker = new CheckpointTimeTracker(ots, config.getStatusTable(), this.streamJob.getStreamId());
        this.checker = new OTSStreamReaderChecker(ots, config);
        this.allShardsMap = OTSHelper.toShardMap(allShards);

        LOG.info("SlaveId: {}, ShardIds: {}, OwnedShards: {}.", slaveId, allShards, ownedShardIds);
        this.ownedShards = new HashMap<String, StreamShard>();
        for (String ownedShardId : ownedShardIds) {
            ownedShards.put(ownedShardId, allShardsMap.get(ownedShardId));
        }

        for (String shardId : this.streamJob.getShardIds()) {
            shardToCheckpointMap.put(shardId, new ShardCheckpoint(shardId, this.streamJob.getVersion(), CheckpointPosition.TRIM_HORIZON, 0));
        }

        findCheckpoints = checker.checkAndSetCheckpoints(checkpointInfoTracker, allShardsMap, streamJob, shardToCheckpointMap);
        if (!findCheckpoints) {
            LOG.info("Checkpoint for stream '{}' in timestamp '{}' is not found.", streamJob.getStreamId(), streamJob.getStartTimeInMillis());
            setWithNearestCheckpoint();
        }

        LOG.info("Find checkpoints: {}.", findCheckpoints);
        for (Map.Entry<String, StreamShard> shard : ownedShards.entrySet()) {
            LOG.info("Shard to process, ShardInfo: [{}], StartCheckpoint: [{}].", shard.getValue(), shardToCheckpointMap.get(shard.getKey()));
        }
        LOG.info("Count of owned shards: {}. ShardIds: {}.", ownedShardIds.size(), ownedShardIds);
    }

    public boolean isFindCheckpoints() {
        return findCheckpoints;
    }

    public Map<String, StreamShard> getAllShardsMap() {
        return allShardsMap;
    }

    public Map<String, StreamShard> getOwnedShards() {
        return ownedShards;
    }

    public Map<String, ShardCheckpoint> getShardToCheckpointMap() {
        return shardToCheckpointMap;
    }

    /**
     * 没有找到上一次任务的checkpoint，需要重新从头开始读。
     * 为了减少扫描的数据量，尝试查找里startTime最近的一次checkpoint。
     */
    private void setWithNearestCheckpoint() {
        long expirationTime = (streamDetails.getExpirationTime() - 1) * TimeUtils.HOUR_IN_MILLIS;
        long timeRangeBegin = System.currentTimeMillis() - expirationTime;
        long timeRangeEnd = this.config.getStartTimestampMillis() - 1;
        if (timeRangeBegin < timeRangeEnd) {
            for (String shardId : ownedShards.keySet()) {
                LOG.info("Try find nearest checkpoint for shard {}, startTime: {}.", shardId, config.getStartTimestampMillis());
                String checkpoint = this.checkpointInfoTracker.getShardLargestCheckpointInTimeRange(shardId, timeRangeBegin, timeRangeEnd);
                if (checkpoint != null) {
                    LOG.info("Found checkpoint for shard {}, checkpoint: {}.", shardId, checkpoint);
                    shardToCheckpointMap.put(shardId, new ShardCheckpoint(shardId, streamJob.getVersion(), checkpoint, 0));
                }
            }
        }
    }

    private int calcThreadPoolSize() {
        int threadNum = 0;
        // 如果配置了thread num，则计算平均每个slave所启动的thread的个数
        if (config.getThreadNum() > 0) {
            threadNum = config.getThreadNum() / slaveNumber.get();
        } else {
            threadNum = Runtime.getRuntime().availableProcessors() * 4 / slaveNumber.get();
        }

        if (threadNum == 0) {
            threadNum = 1;
        }
        LOG.info("ThreadNum: {}.", threadNum);
        return threadNum;
    }

    private Map<String, StreamShard> filterShardsReachEnd(Map<String, StreamShard> ownedShards, Map<String, ShardCheckpoint> allCheckpoints) {
        Map<String, StreamShard> allShardToProcess = new HashMap<String, StreamShard>();
        for (Map.Entry<String, StreamShard> shard : ownedShards.entrySet()) {
            String shardId = shard.getKey();
            if (allCheckpoints.get(shardId).getCheckpoint().equals(CheckpointPosition.SHARD_END)) {
                LOG.info("Shard has reach end, no need to process. ShardId: {}.", shardId);
                // but we need to set checkpoint for this job
                checkpointInfoTracker.writeCheckpoint(streamJob.getEndTimeInMillis(),
                        new ShardCheckpoint(shardId, streamJob.getVersion(), CheckpointPosition.SHARD_END, 0), 0);
            } else {
                allShardToProcess.put(shard.getKey(), shard.getValue());
            }
        }
        return allShardToProcess;
    }

    public void startRead(RecordSender recordSender) {
        int threadPoolSize = calcThreadPoolSize();
        ExecutorService executorService = new ThreadPoolExecutor(
                0, threadPoolSize, 60L, TimeUnit.SECONDS, new ArrayBlockingQueue<Runnable>(ownedShards.size()));
        LOG.info("Start thread pool with size: {}, ShardsCount: {}, SlaveCount: {}.", threadPoolSize, ownedShards.size(), slaveNumber.get());
        try {
            Map<String, StreamShard> allShardToProcess = filterShardsReachEnd(ownedShards, shardToCheckpointMap);
            Map<String, ShardStatusChecker.ProcessState> shardProcessingState = new HashMap<String, ShardStatusChecker.ProcessState>();
            for (String shardId : allShardToProcess.keySet()) {
                shardProcessingState.put(shardId, ShardStatusChecker.ProcessState.BLOCK);
            }

            List<RecordProcessor> processors = new ArrayList<RecordProcessor>();

            // 获取当前所有shard的checkpoint状态，对当前的owned shard执行对应的任务。
            long lastLogTime = System.currentTimeMillis();
            while (!allShardToProcess.isEmpty()) {
                Map<String, ShardCheckpoint> checkpointMap = checkpointInfoTracker.getAllCheckpoints(streamJob.getEndTimeInMillis());

                // 检查当前job的checkpoint，排查是否有其他job误入或者出现不明的shard。
                checkCheckpoint(checkpointMap, streamJob);

                // 找到需要处理的shard以及确定不需要被处理的shard
                List<StreamShard> shardToProcess = new ArrayList<StreamShard>();
                List<StreamShard> shardNoNeedProcess = new ArrayList<StreamShard>();
                List<StreamShard> shardBlocked = new ArrayList<StreamShard>();
                ShardStatusChecker.findShardToProcess(allShardToProcess, allShardsMap, checkpointMap, shardToProcess, shardNoNeedProcess, shardBlocked);

                // 将不需要处理的shard，设置checkpoint，代表本轮处理完毕，且checkpoint为TRIM_HORIZON
                for (StreamShard shard : shardNoNeedProcess) {
                    LOG.info("Skip shard: {}.", shard.getShardId());
                    ShardCheckpoint checkpoint = new ShardCheckpoint(shard.getShardId(), streamJob.getVersion(), CheckpointPosition.TRIM_HORIZON, 0);
                    checkpointInfoTracker.writeCheckpoint(config.getEndTimestampMillis(), checkpoint, 0);
                    shardProcessingState.put(shard.getShardId(), ShardStatusChecker.ProcessState.SKIP);
                }

                for (StreamShard shard : shardToProcess) {
                    RecordProcessor processor = new RecordProcessor(ots, config, streamJob, shard,
                            shardToCheckpointMap.get(shard.getShardId()), !findCheckpoints, checkpointInfoTracker, recordSender);
                    processor.initialize();
                    executorService.submit(processor);
                    processors.add(processor);
                    shardProcessingState.put(shard.getShardId(), ShardStatusChecker.ProcessState.READY);
                }

                // 等待所有任务执行完毕，并且检查每个任务的状态，检查是否发生hang或长时间没有数据
                checkProcessorRunningStatus(processors);

                if (!allShardToProcess.isEmpty()) {
                    TimeUtils.sleepMillis(config.getSlaveLoopInterval());
                }

                long now = System.currentTimeMillis();
                if (now - lastLogTime > config.getSlaveLoggingStatusInterval()) {
                    logShardProcessingState(shardProcessingState);
                    LOG.info("AllCheckpoints: {}", checkpointMap);
                    lastLogTime = now;
                }
            }

            LOG.info("All shard is processing.");
            logShardProcessingState(shardProcessingState);
            // 等待当前分配的shard的读取任务执行完毕后退出。
            while (true) {
                boolean finished = true;
                checkProcessorRunningStatus(processors);
                for (RecordProcessor processor : processors) {
                    RecordProcessor.State state = processor.getState();
                    if (state != RecordProcessor.State.SUCCEED) {
                        LOG.info("Shard is processing, shardId: {}, status: {}.", processor.getShard().getShardId(), state);
                        finished = false;
                    }
                }

                if (finished) {
                    LOG.info("All record processor finished.");
                    break;
                }

                TimeUtils.sleepMillis(config.getSlaveLoopInterval());
            }

        } catch (TableStoreException ex) {
            throw DataXException.asDataXException(new OTSReaderError(ex.getErrorCode(), "SyncClientInterface Error"), ex.toString(), ex);
        } catch (OTSStreamReaderException ex) {
            LOG.error("SlaveId: {}, OwnedShards: {}.", slaveId, ownedShards, ex);
            throw DataXException.asDataXException(OTSReaderError.ERROR, ex.toString(), ex);
        } catch (Exception ex) {
            LOG.error("SlaveId: {}, OwnedShards: {}.", slaveId, ownedShards, ex);
            throw DataXException.asDataXException(OTSReaderError.ERROR, ex.toString(), ex);
        } finally {
            try {
                executorService.shutdownNow();
                executorService.awaitTermination(1, TimeUnit.MINUTES);
            } catch (Exception e) {
                LOG.error("Shutdown encounter exception.", e);
            }
        }
    }

    private void logShardProcessingState(Map<String, ShardStatusChecker.ProcessState> shardProcessingState) {
        StringBuilder sb = new StringBuilder();
        sb.append("Shard running status: \n");
        for (Map.Entry<String, ShardStatusChecker.ProcessState> entry : shardProcessingState.entrySet()) {
            sb.append("ShardId:").append(entry.getKey()).
                    append(", ProcessingState: ").append(entry.getValue()).append("\n");
        }
        LOG.info("Version: {}, Reader status: {}", streamJob.getVersion(), sb.toString());
    }

    private void checkProcessorRunningStatus(List<RecordProcessor> processors) {
        long now = System.currentTimeMillis();
        for (RecordProcessor processor : processors) {
            RecordProcessor.State state = processor.getState();
            StreamShard shard = processor.getShard();
            if (state == RecordProcessor.State.READY || state == RecordProcessor.State.SUCCEED) {
                continue;
            } else if (state == RecordProcessor.State.INTERRUPTED || state == RecordProcessor.State.FAILED) {
                throw new OTSStreamReaderException("Read task for shard '" + shard.getShardId() + "' has failed.");
            } else { // status = RUNNING
                long lastProcessTime = processor.getLastProcessTime();
                if (now - lastProcessTime > OTSStreamReaderConstants.MAX_ONCE_PROCESS_TIME_MILLIS) {
                    throw new OTSStreamReaderException("Process shard timeout, ShardId:" + shard.getShardId() + ", LastProcessTime:"
                            + lastProcessTime + ", MaxProcessTime:" + OTSStreamReaderConstants.MAX_ONCE_PROCESS_TIME_MILLIS + ", Now:" + now + ".");
                }
            }
        }
    }

    void checkCheckpoint(Map<String, ShardCheckpoint> checkpointMap, StreamJob streamJob) {
        for (Map.Entry<String, ShardCheckpoint> entry : checkpointMap.entrySet()) {
            String shardId = entry.getKey();
            String version = entry.getValue().getVersion();
            if (!streamJob.getShardIds().contains(shardId)) {
                LOG.info("Shard '{}' is not found in job. Job: {}.", entry.getKey(), streamJob.getShardIds());
                throw DataXException.asDataXException(OTSReaderError.ERROR, "Some shard from checkpoint is not belong to this job: " + shardId);
            }

            if (!version.equals(streamJob.getVersion())) {
                LOG.info("Version of shard '{}' in checkpoint is not equal with version of this job. " +
                        "Checkpoint version: {}, job version: {}.", shardId, version, streamJob.getVersion());
                throw DataXException.asDataXException(OTSReaderError.ERROR, "Version of checkpoint is not equal with version of this job.");
            }
        }
    }

    public void close() {
        ots.shutdown();
    }
}
