package com.github.msemys.esjc.tcp.handler;

import com.github.msemys.esjc.tcp.TcpCommand;
import com.github.msemys.esjc.tcp.TcpPackage;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.timeout.IdleStateEvent;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Duration;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledFuture;

import static java.util.concurrent.TimeUnit.MILLISECONDS;

public class HeartbeatHandler extends SimpleChannelInboundHandler<TcpPackage> {
    private static final Logger logger = LoggerFactory.getLogger(HeartbeatHandler.class);

    private final long timeoutMillis;
    private Map<UUID, ScheduledFuture<?>> timeoutTasks = new ConcurrentHashMap<>();

    public HeartbeatHandler(Duration timeout) {
        timeoutMillis = timeout.toMillis();
    }

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, TcpPackage msg) throws Exception {
        switch (msg.command) {
            case HeartbeatRequestCommand:
                ctx.writeAndFlush(TcpPackage.newBuilder()
                    .command(TcpCommand.HeartbeatResponseCommand)
                    .correlationId(msg.correlationId)
                    .build());
                break;
            case HeartbeatResponseCommand:
                cancelTimeoutTask(msg.correlationId);
                break;
            default:
                ctx.fireChannelRead(msg);
        }
    }

    @Override
    public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
        cancelAllTimeoutTasks();
    }

    @Override
    public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
        if (evt instanceof IdleStateEvent) {
            UUID correlationId = UUID.randomUUID();
            ScheduledFuture<?> timeoutTask = ctx.executor().schedule(() -> {
                logger.info("Closing TCP connection [{}, L{}] due to HEARTBEAT TIMEOUT.", ctx.channel().remoteAddress(), ctx.channel().localAddress());
                ctx.close();
            }, timeoutMillis, MILLISECONDS);
            timeoutTasks.put(correlationId, timeoutTask);

            ctx.writeAndFlush(TcpPackage.newBuilder()
                .command(TcpCommand.HeartbeatRequestCommand)
                .correlationId(correlationId)
                .build());
        }
    }

    private void cancelTimeoutTask(UUID correlationId) {
        ScheduledFuture<?> timeoutTask = timeoutTasks.remove(correlationId);
        if (timeoutTask != null) {
            timeoutTask.cancel(true);
        }
    }

    private void cancelAllTimeoutTasks() {
        timeoutTasks.forEach((uuid, scheduledFuture) -> scheduledFuture.cancel(true));
    }

}

@@ -384,81 +384,12 @@
 til.
-Map;%0Aimport java.util.UUID;%0Aimport java.util.concurrent.ConcurrentHashMap
+UUID
 ;%0Aim
@@ -716,17 +716,16 @@
 ate 
-Map%3CUUID,
+volatile
  Sch
@@ -739,17 +739,16 @@
 uture%3C?%3E
-%3E
  timeout
@@ -755,37 +755,8 @@
 Task
-s = new ConcurrentHashMap%3C%3E()
 ;%0A%0A 
@@ -1343,33 +1343,16 @@
 outTask(
-msg.correlationId
 );%0A     
@@ -1550,27 +1550,24 @@
       cancel
-All
 TimeoutTasks
@@ -1565,17 +1565,16 @@
 eoutTask
-s
 ();%0A    
@@ -1740,21 +1740,136 @@
         
-UUID 
+ctx.writeAndFlush(TcpPackage.newBuilder()%0A                .command(TcpCommand.HeartbeatRequestCommand)%0A                .
 correlat
@@ -1873,19 +1873,17 @@
 lationId
- = 
+(
 UUID.ran
@@ -1891,18 +1891,21 @@
 omUUID()
-;%0A
+)%0A   
         
@@ -1908,34 +1908,39 @@
         
-ScheduledFuture%3C?%3E
+ .build());%0A           
  timeout
@@ -2216,397 +2216,55 @@
     
-    timeoutTasks.put(correlationId, timeoutTask);%0A%0A            ctx.writeAndFlush(TcpPackage.newBuilder()%0A                .command(TcpCommand.HeartbeatRequestCommand)%0A                .correlationId(correlationId)%0A                .build());%0A        %7D%0A    %7D%0A%0A    private void cancelTimeoutTask(UUID correlationId) %7B%0A        ScheduledFuture%3C?%3E timeoutTask = timeoutTasks.remove(correlationId);
+%7D%0A    %7D%0A%0A    private void cancelTimeoutTask() %7B
 %0A   
@@ -2345,44 +2345,13 @@
     
-%7D%0A
     
-%7D%0A%0A    private void cancelAllT
+t
 imeo
@@ -2360,13 +2360,16 @@
 Task
-s() %7B
+ = null;
 %0A   
@@ -2377,86 +2377,9 @@
     
-timeoutTasks.forEach((uuid, scheduledFuture) -%3E scheduledFuture.cancel(true));
+%7D
 %0A   

Revert "! [Heartbeats] Fixed handling of multiple heartbeats"