/*
 * Copyright (C) 2016-2022 ActionTech.
 * based on code by MyCATCopyrightHolder Copyright (c) 2013, OpenCloudDB/MyCAT.
 * License: http://www.gnu.org/licenses/gpl.html GPL version 2 or higher.
 */
package com.actiontech.dble.backend.mysql.nio.handler;

import com.actiontech.dble.backend.pool.PooledConnectionListener;
import com.actiontech.dble.backend.pool.util.TimerHolder;
import com.actiontech.dble.net.connection.BackendConnection;
import com.actiontech.dble.net.connection.PooledConnection;
import com.actiontech.dble.net.mysql.FieldPacket;
import com.actiontech.dble.net.mysql.PingPacket;
import com.actiontech.dble.net.mysql.RowDataPacket;
import com.actiontech.dble.net.service.AbstractService;
import com.actiontech.dble.net.service.WriteFlags;
import io.netty.util.Timeout;
import org.jetbrains.annotations.NotNull;

import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * heartbeat check for mysql connections
 *
 * @author wuzhih
 */
public class ConnectionHeartBeatHandler implements ResponseHandler {

    private final Object heartbeatLock;
    private volatile Timeout heartbeatTimeout;
    private final BackendConnection conn;
    private final PooledConnectionListener listener;
    private boolean returned = false;
    private boolean finished = false;

    public ConnectionHeartBeatHandler(BackendConnection conn, boolean isBlock, PooledConnectionListener listener) {
        conn.getBackendService().setResponseHandler(this);
        this.conn = conn;
        this.listener = listener;
        if (isBlock) {
            this.heartbeatLock = new Object();
        } else {
            this.heartbeatLock = null;
        }
    }

    public boolean ping(long timeout) {
        if (heartbeatLock != null) {
            final long deadline = System.currentTimeMillis() + timeout;
            synchronized (heartbeatLock) {
                conn.getService().write(PingPacket.PING, WriteFlags.QUERY_END);
                try {
                    while (!returned) {
                        timeout = deadline - System.currentTimeMillis();
                        if (timeout <= 0L) {
                            returned = true;
                        } else {
                            heartbeatLock.wait(timeout);
                        }
                    }
                } catch (InterruptedException e) {
                    returned = true;
                }
            }
        } else {
            heartbeatTimeout = TimerHolder.getTimer().newTimeout(timeout1 -> conn.businessClose("conn heart timeout"), timeout, TimeUnit.MILLISECONDS);
            conn.getService().write(PingPacket.PING, WriteFlags.QUERY_END);
        }

        return finished;
    }

    /**
     * if the query returns ok than just release the connection
     * and go on check the next one
     *
     * @param ok
     * @param service
     */
    @Override
    public void okResponse(byte[] ok, @NotNull AbstractService service) {
        if (heartbeatLock != null) {
            synchronized (heartbeatLock) {
                if (!returned) {
                    returned = true;
                    finished = true;
                    heartbeatLock.notifyAll();
                }
            }
            return;
        }

        heartbeatTimeout.cancel();
        listener.onHeartbeatSuccess((PooledConnection) service.getConnection());
    }

    @Override
    public void connectionClose(@NotNull AbstractService service, String reason) {
        if (heartbeatLock != null) {
            synchronized (heartbeatLock) {
                if (!returned) {
                    returned = true;
                    finished = false;
                    heartbeatLock.notifyAll();
                }
            }
        }
    }

    @Override
    public void fieldEofResponse(byte[] header, List<byte[]> fields, List<FieldPacket> fieldPackets, byte[] eof, boolean isLeft, @NotNull AbstractService service) {

    }

    @Override
    public boolean rowResponse(byte[] rowNull, RowDataPacket rowPacket, boolean isLeft, @NotNull AbstractService service) {
        return false;
    }

    @Override
    public void rowEofResponse(byte[] eof, boolean isLeft, @NotNull AbstractService service) {

    }

    @Override
    public void connectionError(Throwable e, Object attachment) {
        // not called
    }

    @Override
    public void connectionAcquired(com.actiontech.dble.net.connection.BackendConnection connection) {

    }

    @Override
    public void errorResponse(byte[] err, @NotNull AbstractService service) {

    }
}
