/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.eclipse.aether.transport.jetty;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.nio.ByteBuffer;
import java.nio.channels.ByteChannel;
import java.nio.channels.Channels;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.ReadableByteChannel;
import java.nio.channels.SeekableByteChannel;
import java.nio.file.Files;
import java.nio.file.StandardOpenOption;
import java.util.Objects;
import java.util.function.Supplier;

import org.eclipse.aether.spi.connector.transport.PutTask;
import org.eclipse.jetty.client.ByteBufferRequestContent;
import org.eclipse.jetty.client.Request;
import org.eclipse.jetty.io.ByteBufferPool;
import org.eclipse.jetty.io.Content;
import org.eclipse.jetty.io.RetainableByteBuffer;
import org.eclipse.jetty.io.internal.ByteChannelContentSource;
import org.eclipse.jetty.util.BufferUtil;
import org.eclipse.jetty.util.ExceptionUtil;
import org.eclipse.jetty.util.IO;
import org.eclipse.jetty.util.TypeUtil;
import org.eclipse.jetty.util.thread.AutoLock;
import org.eclipse.jetty.util.thread.SerializedInvoker;

/**
 * Heavily inspired by Jetty's {@code org.eclipse.jetty.io.internal.ByteChannelContentSource} but adjusted to deal with
 * {@link ReadableByteChannel}s and to support rewind (to be able to retry the requests).
 * Also, Jetty's {@code ByteChannelContentSource} is an internal package so should not be used directly.
 * @see <a href="https://javadoc.jetty.org/jetty-12/org/eclipse/jetty/io/internal/ByteChannelContentSource.html">ByteChannelContentSource</a>
 * @see <a href="https://github.com/jetty/jetty.project/issues/14324">Jetty Issue #14324</a>
 */
public class PutTaskRequestContent extends ByteBufferRequestContent implements Request.Content {

    public static Request.Content from(PutTask putTask) {
        Supplier<ReadableByteChannel> newChannelSupplier;
        if (putTask.getDataPath() != null) {
            newChannelSupplier = () -> {
                try {
                    return Files.newByteChannel(putTask.getDataPath(), StandardOpenOption.READ);
                } catch (IOException e) {
                    throw new UncheckedIOException(e);
                }
            };
        } else {
            newChannelSupplier = () -> {
                try {
                    return Channels.newChannel(putTask.newInputStream());
                } catch (IOException e) {
                    throw new UncheckedIOException(e);
                }
            };
        }
        return new PutTaskRequestContent(null, newChannelSupplier, 0L, putTask.getDataLength());
    }

    private final AutoLock lock = new AutoLock();
    private final SerializedInvoker invoker = new SerializedInvoker(ByteChannelContentSource.class);
    private final ByteBufferPool.Sized byteBufferPool;
    private ReadableByteChannel byteChannel;
    private final long offset;
    private final long length;
    private RetainableByteBuffer buffer;
    private long offsetRemaining;
    private long totalRead;
    private Runnable demandCallback;
    private Content.Chunk terminal;
    /** Only necessary for rewind support when leveraging the input stream. */
    private final Supplier<ReadableByteChannel> newByteChannelSupplier;

    /**
     * Create a {@link ByteChannelContentSource} which reads from a {@link ByteChannel}.
     * @param byteBufferPool The {@link org.eclipse.jetty.io.ByteBufferPool.Sized} to use for any internal buffers.
     * @param newByteChannelSupplier The {@link ByteChannel} supplier.
     */
    protected PutTaskRequestContent(
            ByteBufferPool.Sized byteBufferPool, Supplier<ReadableByteChannel> newByteChannelSupplier) {
        this(byteBufferPool, newByteChannelSupplier, 0L, -1L);
    }

    /**
     * Create a {@link ByteChannelContentSource} which reads from a {@link ByteChannel}.
     * If the {@link ByteChannel} is an instance of {@link SeekableByteChannel} the implementation will use
     * {@link SeekableByteChannel#position(long)} to navigate to the starting offset.
     * @param byteBufferPool The {@link org.eclipse.jetty.io.ByteBufferPool.Sized} to use for any internal buffers.
     * @param newByteChannelSupplier The {@link ByteChannel} supplier.
     * @param offset the offset byte of the content to start from.
     *               Must be greater than or equal to 0 and less than the content length (if known).
     * @param length the length of the content to make available, -1 for the full length.
     *               If the size of the content is known, the length may be truncated to the content size minus the offset.
     * @throws IndexOutOfBoundsException if the offset or length are out of range.
     * @see TypeUtil#checkOffsetLengthSize(long, long, long)
     */
    protected PutTaskRequestContent(
            ByteBufferPool.Sized byteBufferPool,
            Supplier<ReadableByteChannel> newByteChannelSupplier,
            long offset,
            long length) {
        this.byteBufferPool = Objects.requireNonNullElse(byteBufferPool, ByteBufferPool.SIZED_NON_POOLING);
        this.byteChannel = newByteChannelSupplier.get();
        this.offset = offset;
        this.length = TypeUtil.checkOffsetLengthSize(offset, length, -1L);
        offsetRemaining = offset;
        this.newByteChannelSupplier = newByteChannelSupplier;
    }

    protected ReadableByteChannel open() throws IOException {
        return byteChannel;
    }

    @Override
    public void demand(Runnable demandCallback) {
        try (AutoLock ignored = lock.lock()) {
            if (this.demandCallback != null) {
                throw new IllegalStateException("demand pending");
            }
            this.demandCallback = demandCallback;
        }
        invoker.run(this::invokeDemandCallback);
    }

    private void invokeDemandCallback() {
        Runnable demandCallback;
        try (AutoLock ignored = lock.lock()) {
            demandCallback = this.demandCallback;
            this.demandCallback = null;
        }
        if (demandCallback != null) {
            ExceptionUtil.run(demandCallback, this::fail);
        }
    }

    protected void lockedSetTerminal(Content.Chunk terminal) {
        assert lock.isHeldByCurrentThread();
        if (terminal != null) {
            ExceptionUtil.addSuppressedIfNotAssociated(terminal.getFailure(), terminal.getFailure());
        }
        IO.close(byteChannel);
        if (buffer != null) {
            buffer.release();
        }
        buffer = null;
    }

    private void lockedEnsureOpenOrTerminal() {
        assert lock.isHeldByCurrentThread();
        if (terminal == null && (byteChannel == null || !byteChannel.isOpen())) {
            try {
                byteChannel = open();
                if (byteChannel == null || !byteChannel.isOpen()) {
                    lockedSetTerminal(Content.Chunk.from(new ClosedChannelException(), true));
                } else if (byteChannel instanceof SeekableByteChannel) {
                    ((SeekableByteChannel) byteChannel).position(offset);
                    offsetRemaining = 0;
                }
            } catch (IOException e) {
                lockedSetTerminal(Content.Chunk.from(e, true));
            }
        }
    }

    @Override
    public Content.Chunk read() {
        try (AutoLock ignored = lock.lock()) {
            lockedEnsureOpenOrTerminal();

            if (terminal != null) {
                return terminal;
            }

            if (length == 0) {
                lockedSetTerminal(Content.Chunk.EOF);
                return Content.Chunk.EOF;
            }

            if (buffer == null) {
                buffer = byteBufferPool.acquire();
            } else if (buffer.isRetained()) {
                buffer.release();
                buffer = byteBufferPool.acquire();
            }

            try {
                ByteBuffer byteBuffer = buffer.getByteBuffer();
                if (offsetRemaining > 0) {
                    // Discard all bytes read until we reach the staring offset.
                    while (offsetRemaining > 0) {
                        BufferUtil.clearToFill(byteBuffer);
                        byteBuffer.limit((int) Math.min(buffer.capacity(), offsetRemaining));
                        int read = byteChannel.read(byteBuffer);
                        if (read < 0) {
                            lockedSetTerminal(Content.Chunk.EOF);
                            return terminal;
                        }
                        if (read == 0) {
                            return null;
                        }

                        offsetRemaining -= read;
                    }
                }

                BufferUtil.clearToFill(byteBuffer);
                if (length > 0) {
                    byteBuffer.limit((int) Math.min(buffer.capacity(), length - totalRead));
                }
                int read = byteChannel.read(byteBuffer);
                BufferUtil.flipToFlush(byteBuffer, 0);
                if (read == 0) {
                    return null;
                }
                if (read > 0) {
                    totalRead += read;
                    buffer.retain();
                    if (length < 0 || totalRead < length) {
                        return Content.Chunk.asChunk(byteBuffer, false, buffer);
                    }

                    Content.Chunk last = Content.Chunk.asChunk(byteBuffer, true, buffer);
                    lockedSetTerminal(Content.Chunk.EOF);
                    return last;
                }
                lockedSetTerminal(Content.Chunk.EOF);
            } catch (Throwable t) {
                lockedSetTerminal(Content.Chunk.from(t, true));
            }
        }
        return terminal;
    }

    @Override
    public void fail(Throwable failure) {
        try (AutoLock ignored = lock.lock()) {
            lockedSetTerminal(Content.Chunk.from(failure, true));
        }
    }

    @Override
    public long getLength() {
        return length;
    }

    @Override
    public boolean rewind() {
        try (AutoLock ignored = lock.lock()) {
            // open a new ByteChannel if we don't have a SeekableByteChannel.
            if (!(byteChannel instanceof SeekableByteChannel)) {
                try {
                    byteChannel.close();
                } catch (IOException e) {
                    throw new UncheckedIOException(e);
                }
                byteChannel = newByteChannelSupplier.get();
                offsetRemaining = 0;
                totalRead = 0;
                return true;
            }

            // We can remove terminal condition for a rewind that is likely to occur
            if (terminal != null
                    && !Content.Chunk.isFailure(terminal)
                    && (byteChannel == null || byteChannel instanceof SeekableByteChannel)) {
                terminal = null;
            }

            lockedEnsureOpenOrTerminal();
            if (terminal != null || byteChannel == null || !byteChannel.isOpen()) {
                return false;
            }

            try {
                ((SeekableByteChannel) byteChannel).position(offset);
                offsetRemaining = 0;
                totalRead = 0;
                return true;
            } catch (Throwable t) {
                lockedSetTerminal(Content.Chunk.from(t, true));
            }

            return true;
        }
    }
}
