/*
 * Copyright (C) 2016-2022 ActionTech.
 * License: http://www.gnu.org/licenses/gpl.html GPL version 2 or higher.
 */

package com.actiontech.dble.net.impl.aio;

import com.actiontech.dble.net.SocketWR;
import com.actiontech.dble.net.WriteOutTask;
import com.actiontech.dble.net.connection.AbstractConnection;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousSocketChannel;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicBoolean;

public class AIOSocketWR extends SocketWR {
    protected static final Logger LOGGER = LoggerFactory.getLogger(SocketWR.class);
    private static final AIOReadHandler AIO_READ_HANDLER = new AIOReadHandler();
    private static final AIOWriteHandler AIO_WRITE_HANDLER = new AIOWriteHandler();
    private AsynchronousSocketChannel channel;
    protected AbstractConnection con;
    protected final AtomicBoolean writing = new AtomicBoolean(false);
    private ConcurrentLinkedQueue<WriteOutTask> writeQueue;

    private volatile WriteOutTask leftoverWriteTask;
    volatile boolean writeDataErr = false;

    public AIOSocketWR() {

    }

    public void initFromConnection(AbstractConnection connection) {
        this.con = connection;
        this.channel = (AsynchronousSocketChannel) connection.getChannel();
        this.writeQueue = connection.getWriteQueue();
    }

    @Override
    public void asyncRead() throws IOException {
        if (con.isClosed()) {
            throw new IOException("read from closed channel cause error");
        }
        try {
            ByteBuffer theBuffer = con.findReadBuffer();
            if (theBuffer.hasRemaining()) {
                channel.read(theBuffer, this, AIO_READ_HANDLER);
            } else {
                throw new IllegalArgumentException("full buffer to read ");
            }
        } finally {
            //prevent  asyncClose and read operation happened Concurrently.
            if (con.isClosed() && con.getBottomReadBuffer() != null) {
                con.recycleReadBuffer();
            }
        }

    }

    private void asyncWrite(final ByteBuffer buffer) {
        buffer.flip();
        this.channel.write(buffer, this, AIO_WRITE_HANDLER);

    }

    @Override
    public boolean isWriteComplete() {
        if (writeDataErr) {
            return true;
        }
        int count = 0;
        while (!writing.compareAndSet(false, true) && count < 5) {
            try {
                Thread.sleep(10);
                count++;
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
        if (count >= 5) {
            return false;
        }
        try {
            ByteBuffer buffer = leftoverWriteTask == null ? null : leftoverWriteTask.getBuffer();
            if (buffer != null && buffer.hasRemaining()) {
                return false;
            }
            return writeQueue.isEmpty();

        } finally {
            writing.set(false);
        }

    }

    /**
     * return true ,means no more data
     *
     * @return
     */
    private boolean write0() {
        if (!writing.compareAndSet(false, true)) {
            return false;
        }
        ByteBuffer theBuffer = leftoverWriteTask == null ? null : leftoverWriteTask.getBuffer();
        if (theBuffer == null || !theBuffer.hasRemaining()) {
            if (theBuffer != null) {
                con.recycle(theBuffer);
                leftoverWriteTask = null;
            }
            // poll again
            WriteOutTask task = writeQueue.poll();
            // more data
            if (task != null) {
                ByteBuffer buffer = task.getBuffer();
                boolean quitFlag = task.closeFlag();
                if (buffer.limit() == 0) {
                    con.recycle(buffer);
                    leftoverWriteTask = null;
                    con.close("quit cmd");
                    writing.set(false);
                    return true;
                } else {
                    leftoverWriteTask = task;
                    asyncWrite(buffer);
                    if (quitFlag) {
                        con.recycle(buffer);
                        con.close(con.getCloseReason());
                        return true;
                    }
                    return false;
                }
            } else {
                // no buffer
                writing.set(false);
                return true;
            }
        } else {
            boolean quitFlag = leftoverWriteTask == null ? false : leftoverWriteTask.closeFlag();
            theBuffer.compact();
            asyncWrite(theBuffer);
            if (quitFlag) {
                con.recycle(theBuffer);
                con.close(con.getCloseReason());
                return true;
            }
            return false;
        }

    }

    protected void onWriteFinished(int result) {
        con.writeStatistics(result);
        boolean noMoreData = this.write0();
        if (noMoreData) {
            this.doNextWriteCheck();
        }

    }

    public void doNextWriteCheck() {

        boolean noMoreData = false;
        noMoreData = this.write0();
        if (noMoreData && !con.getWriteQueue().isEmpty()) {
            this.write0();
        }
    }

    public boolean registerWrite(ByteBuffer buffer) {
        leftoverWriteTask = new WriteOutTask(buffer, false);
        buffer.flip();
        try {
            write0();
        } catch (Exception e) {
            //SLB no exception in AIO
            if (LOGGER.isDebugEnabled()) {
                LOGGER.debug("caught err:", e);
            }
            con.close("err:" + e);
            return false;
        }
        return true;
    }

    @Override
    public void disableRead() {

    }

    @Override
    public void enableRead() {

    }

    @Override
    public void disableReadForever() {
        try {
            channel.shutdownInput();
        } catch (IOException e) {
            LOGGER.info("disable read error", e);
        }
    }

    @Override
    public void shutdownInput() throws IOException {
        channel.shutdownInput();
    }

    @Override
    public void closeSocket() throws IOException {
        channel.close();
    }

    public void setWriteDataErr(boolean writeDataErr) {
        this.writeDataErr = writeDataErr;
    }

    @Override
    public boolean canNotWrite() {
        return writeDataErr;
    }
}


