package window;

import org.apache.flink.runtime.state.CheckpointListener;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.java.typeutils.ListTypeInfo;
import org.apache.flink.api.java.typeutils.MapTypeInfo;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Consumer;

public class EntityKeyedProcessFunction
        extends KeyedProcessFunction<String, DataEntity, DataEntity>
        implements CheckpointedFunction, CheckpointListener {

    static final class SubState implements Serializable,Iterable<Map.Entry<String, List<DataEntity>>>{
        Map<String, List<DataEntity>> state = new HashMap<>();

        public Set<Map.Entry<String, List<DataEntity>>> entrySet () {
            return state.entrySet();
        }

        public Set<String> keySet(){
            return state.keySet();
        }

        public boolean keyExist(String key){
            return state.containsKey(key);
        }

        public void addValue(String key, DataEntity entity){
            if (keyExist(key)){
                state.get(key).add(entity);
            }else {
                List<DataEntity> ld = new ArrayList<>();
                ld.add(entity);
                state.put(key,ld);
            }
        }

        public Map<String, List<DataEntity>> getState() {
            return state;
        }

        @Override
        public Iterator<Map.Entry<String, List<DataEntity>>> iterator() {
            return state.entrySet().iterator();
        }

        @Override
        public void forEach(Consumer<? super Map.Entry<String, List<DataEntity>>> action) {
            this.state.entrySet().forEach(action);
        }
    }


    protected static final Logger log = LoggerFactory.getLogger(EntityKeyedProcessFunction.class);


    /**
     * key, the minute of the time
     * <p>
     * map key, the name of the entity
     * map value, the entity of the key
     */
    private MapStateDescriptor<String, SubState> entityStatesDesc =
            new MapStateDescriptor<String, SubState>("entityMap",
                    String.class,
                    SubState.class);

    private transient MapState<String, SubState> entityStates;
    private transient Map<Long, SubState> state;

    @Override
    public void snapshotState(FunctionSnapshotContext context) throws Exception {
        log.info("begin to submit checkpoint {}", context.getCheckpointId());
        this.entityStates.clear();

        for(Map.Entry<Long,SubState> entry : this.state.entrySet()){
            this.entityStates.put(entry.getKey()+"" ,entry.getValue());
            log.info("about to checkpoint on key={}",entry.getKey());
        }
    }

    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {
        this.entityStates = context.getKeyedStateStore().getMapState(entityStatesDesc);

        if (context.isRestored()){
            printStates();
        }

        log.info("initializeState in {}", this.getClass());
    }

    private void printStates() throws Exception{
        log.info("**************begin to print restored states.");
        for(Map.Entry<String,SubState> entry : this.entityStates.entries()){
            log.info(">>>>>>>>>the key is {}", entry.getKey());
            SubState map = entry.getValue();
            for(Map.Entry<String,List<DataEntity>> me : map.entrySet()){
                log.info("the value is name = {}, value = {}", me.getKey(), me.getValue());
            }
        }
        log.info("**************end to print restored states.");
    }

    @Override
    public void notifyCheckpointComplete(long checkpointId) throws Exception {
        log.info("Checkpoint {} has completed.", checkpointId);
        printStates();
    }

    @Override
    public void notifyCheckpointAborted(long checkpointId) throws Exception {
        log.info("Checkpoint {} has Aborted.", checkpointId);
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        log.info("open in {}", this.getClass());

        state = new HashMap<>();

        for(Map.Entry<String,SubState> entry : this.entityStates.entries()){
            log.info(">>>>>>>>>added the key is {}", entry.getKey());
            try {
                state.put(Long.parseLong(entry.getKey()),entry.getValue());
            }catch (NumberFormatException nfe){
                log.info("Fail to parse long from {}",entry.getKey(),nfe);
            }

        }
    }

    @Override
    public void processElement(DataEntity dataEntity, Context context, Collector<DataEntity> collector) throws Exception {
        log.info("Receive data whose key as {}", dataEntity.getName());
        long now = System.currentTimeMillis();

        Long minute = roundToMinute(now);
        SubState map = this.state.get(minute);
        if(map == null){
            map = new SubState();

        }

        map.addValue(dataEntity.getName(),dataEntity);

        this.state.put(minute,map);

        long addOne = getOneMinute(now);

        context.timerService().registerProcessingTimeTimer(addOne);
    }

    private long roundToMinute(long time){
        long mi = time/1000/60;
        long nmi = mi * 1000 * 60;
        return nmi;
    }

    private Long getOneMinute(long timestamp){
        long now = roundToMinute(timestamp);
        return now + 1000*60;
    }

    private Long getThreeMinute(long timestamp){
        long now = roundToMinute(timestamp);
        return now - 1000*60*3;
    }

    @Override
    public void onTimer(long timestamp, OnTimerContext ctx, Collector<DataEntity> out) throws Exception {

        log.info("On Timer minute {} on key {}", timestamp, ctx.getCurrentKey());
        long ti = getThreeMinute(timestamp);
        Iterator<Map.Entry<Long, SubState>> itl = this.state.entrySet().iterator();
        while (itl.hasNext()){
            Map.Entry<Long, SubState> entry = itl.next();
            if (entry.getKey() <= ti){
                log.info("remove data on key {} and time {}",entry.getValue().keySet(),entry.getKey());
                itl.remove();
            }
        }
    }
}