package cn.axzo.msg.center.utils;

import com.baomidou.mybatisplus.core.toolkit.support.SFunction;
import com.baomidou.mybatisplus.extension.conditions.query.LambdaQueryChainWrapper;
import org.jetbrains.annotations.NotNull;

import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.function.Supplier;

/**
 * @author yanglin
 */
public class RecordCursor<T> implements Iterable<List<T>> {

    private final int batchSize;
    private final SFunction<T, Long> idGetter;
    private final Supplier<LambdaQueryChainWrapper<T>> query;

    public RecordCursor(SFunction<T, Long> idGetter,
                        Supplier<LambdaQueryChainWrapper<T>> query) {
        this(1000, idGetter, query);
    }

    public RecordCursor(int batchSize,
                        SFunction<T, Long> idGetter,
                        Supplier<LambdaQueryChainWrapper<T>> query) {
        this.batchSize = batchSize;
        this.idGetter = idGetter;
        this.query = query;
    }

    @Override @NotNull
    public Iterator<List<T>> iterator() {
        return new Iterator<List<T>>() {

            long maxId = 0;
            // single thread producer
            volatile List<T> batch;
            Boolean batchConsumed;

            @Override
            public boolean hasNext() {
                if (batchConsumed == null || batchConsumed) {
                    batch = nextBatch();
                    batchConsumed = false;
                }
                return !batch.isEmpty();
            }

            @Override
            public List<T> next() {
                if (!hasNext())
                    throw new NoSuchElementException();
                batchConsumed = true;
                return batch;
            }

            private List<T> nextBatch() {
                List<T> records = query.get()
                        .gt(idGetter, maxId)
                        .orderByAsc(idGetter)
                        .last("LIMIT " + batchSize)
                        .list();
                if (!records.isEmpty()) {
                    T maxIdRecord = records.get(records.size() - 1);
                    maxId = idGetter.apply(maxIdRecord);
                }
                return records;
            }
        };
    }
}