/**
 * Copyright (C) 2025, Leo Galambos
 * All rights reserved.
 * 
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 * 
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 * 
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * 3. All advertising materials mentioning features or use of this software must
 *    display the following acknowledgement:
 *    This product includes software developed by the Egothor project.
 *
 * 4. Neither the name of the copyright holder nor the names of its contributors
 *    may be used to endorse or promote products derived from this software
 *    without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */
package zeroecho.util;

import java.io.IOException;
import java.io.InputStream;
import java.security.InvalidParameterException;

import org.bouncycastle.crypto.BufferedBlockCipher;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.DataLengthException;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.StreamCipher;
import org.bouncycastle.crypto.modes.AEADBlockCipher;

import zeroecho.util.aes.AesCipherType;
import zeroecho.util.aes.AesParameters;
import zeroecho.util.aes.AesSupport;

/**
 * A builder for creating {@link InputStream}s that apply symmetric encryption
 * or decryption using a variety of cipher types (block, stream, or AEAD).
 * <p>
 * Supports Bouncy Castle's {@link BufferedBlockCipher}, {@link StreamCipher},
 * and {@link AEADBlockCipher} implementations. This builder validates
 * configuration and initializes the cipher before wrapping the input stream in
 * a cipher-specific processing stream.
 */
public final class SymmetricStreamBuilder {
    /** Internal buffer size used during stream processing. */
    public static final int BUF_SIZE = 4096;

    private InputStream in;
    private Object cipher;
    private CipherParameters params;

    private SymmetricStreamBuilder() {
        // private constructor
    }

    /**
     * Creates a new instance of {@code SymmetricStreamBuilder}.
     *
     * @return a new builder instance
     */
    public static SymmetricStreamBuilder newBuilder() {
        return new SymmetricStreamBuilder();
    }

    /**
     * Sets the input stream that will be encrypted or decrypted.
     *
     * @param in the source {@link InputStream}
     * @return this builder instance
     */
    public SymmetricStreamBuilder withInputStream(final InputStream in) {
        this.in = in;
        return this;
    }

    /**
     * Configures the cipher to be used for symmetric encryption or decryption.
     *
     * This method accepts either a predefined {@link AesCipherType}, or a concrete
     * cipher instance such as:
     * <ul>
     * <li>{@link org.bouncycastle.crypto.StreamCipher}</li>
     * <li>{@link org.bouncycastle.crypto.BufferedBlockCipher}</li>
     * <li>{@link org.bouncycastle.crypto.modes.AEADBlockCipher}</li>
     * </ul>
     * If an {@link AesCipherType} is provided, it is internally converted into the
     * appropriate cipher instance via its {@code createCipher()} method.
     *
     * Any unsupported cipher type will result in an
     * {@link InvalidParameterException}.
     *
     * @param cipher the cipher type or instance to use for stream processing; must
     *               be an instance of {@code AesCipherType}, {@code StreamCipher},
     *               {@code BufferedBlockCipher}, or {@code AEADBlockCipher}
     *
     * @return this {@code SymmetricStreamBuilder} instance for fluent chaining
     *
     * @throws InvalidParameterException if the provided cipher is of an unsupported
     *                                   type
     *
     * @see AesCipherType
     * @see org.bouncycastle.crypto.StreamCipher
     * @see org.bouncycastle.crypto.BufferedBlockCipher
     * @see org.bouncycastle.crypto.modes.AEADBlockCipher
     */
    public SymmetricStreamBuilder withCipher(final Object cipher) {
        this.cipher = switch (cipher) {
            case AesCipherType aes -> aes.createCipher();
            case BufferedBlockCipher buff -> buff;
            case StreamCipher stream -> stream;
            case AEADBlockCipher aead -> aead;
            default -> throw new InvalidParameterException("Unsupported cipher type: " + cipher.getClass().getName());
        };
        return this;
    }

    /**
     * Sets the cipher parameters (e.g., key, IV, nonce).
     *
     * @param params the cipher parameters
     * @return this builder instance
     */
    public SymmetricStreamBuilder withParameters(final CipherParameters params) {
        this.params = params;
        return this;
    }

    /**
     * Configures this {@link SymmetricStreamBuilder} with both the AES cipher
     * instance and its associated parameters, derived from the specified
     * {@link AesParameters}.
     *
     * <p>
     * This method performs two key operations:
     * </p>
     * <ul>
     * <li>Initializes the cipher by invoking {@link AesCipherType#createCipher()}
     * on {@code params.cipherType()}.</li>
     * <li>Generates the corresponding {@link CipherParameters} using the provided
     * secret key, initialization vector (IV), and optional Additional Authenticated
     * Data (AAD).</li>
     * </ul>
     *
     * <p>
     * If the selected cipher type supports AEAD (e.g., AES-GCM), the
     * {@code params.aad()} value will be incorporated into the parameter set. For
     * non-AEAD modes, any provided AAD is ignored.
     * </p>
     *
     * @param params the {@link AesParameters} object containing the cipher type,
     *               secret key, IV, and optional AAD; must not be {@code null}
     * @return this {@code SymmetricStreamBuilder} instance for method chaining
     * @throws NullPointerException if {@code params} is {@code null}
     */
    public SymmetricStreamBuilder withCipherAndParameters(final AesParameters params) {
        withCipher(params.cipherType().createCipher());
        this.params = AesSupport.getParameters(params.cipherType(), params.key(), params.iv(), params.aad());
        return this;
    }

    /**
     * Builds a stream that encrypts data from the configured input stream.
     *
     * @return an {@link InputStream} that provides encrypted output
     * @throws IllegalStateException if the builder is not fully configured
     */
    public InputStream buildEncryptingStream() {
        return buildStream(true);
    }

    /**
     * Builds a stream that decrypts data from the configured input stream.
     *
     * @return an {@link InputStream} that provides decrypted output
     * @throws IllegalStateException if the builder is not fully configured
     */
    public InputStream buildDecryptingStream() {
        return buildStream(false);
    }

    private InputStream buildStream(final boolean forEncryption) {
        if (cipher == null || in == null || params == null) {
            throw new IllegalStateException("InputStream, cipher, and parameters must be provided.");
        }

        return switch (cipher) {
            case BufferedBlockCipher bbc -> {
                bbc.init(forEncryption, params);
                yield new BufferedBlockCipherInputStream(in, bbc);
            }
            case AEADBlockCipher aead -> {
                aead.init(forEncryption, params);
                yield new AEADBlockCipherInputStream(in, aead);
            }
            case StreamCipher sc -> {
                sc.init(forEncryption, params);
                yield new StreamCipherInputStream(in, sc);
            }
            default -> throw new IllegalStateException("Unsupported cipher type: " + cipher.getClass().getName());
        };
    }
}

/**
 * Abstract base class for cipher-based {@link InputStream} implementations that
 * process encrypted or decrypted data using a symmetric cipher.
 *
 * Manages buffered input/output handling, stream reading, and stream
 * finalization logic. Concrete subclasses implement cipher-specific processing
 * through the {@link #processCipher} method.
 */
abstract class SymmetricCipherInputStreamBase extends InputStream {
    /**
     * The underlying input stream supplying raw (encrypted or plaintext) data.
     */
    protected final InputStream in;
    /**
     * Internal buffer used to read data from the underlying input stream.
     */
    protected final byte[] inputBuffer = new byte[SymmetricStreamBuilder.BUF_SIZE];
    /**
     * Buffer holding the output from cipher processing.
     */
    protected final byte[] outputBuffer;
    /**
     * Current read position within the output buffer.
     */
    protected int outputPos;
    /**
     * Total number of valid bytes currently in the output buffer.
     */
    protected int outputLen;
    /**
     * Indicates whether the cipher has been finalized (no more data to process).
     */
    protected boolean finalProcessed;

    /**
     * Constructs a cipher input stream with the specified underlying input and
     * output buffer size.
     *
     * @param in               the underlying input stream
     * @param outputBufferSize the size of the buffer to hold processed output
     */
    protected SymmetricCipherInputStreamBase(final InputStream in, final int outputBufferSize) {
        super();

        this.in = in;
        this.outputBuffer = new byte[outputBufferSize];
    }

    /**
     * Reads a single byte from the encrypted or decrypted stream.
     *
     * @return the byte read, or {@code -1} if end of stream
     * @throws IOException if an I/O or cipher processing error occurs
     */
    @Override
    public int read() throws IOException {
        if (outputPos >= outputLen && !fillBuffer()) {
            return -1;
        }
        return outputBuffer[outputPos++] & 0xFF;
    }

    /**
     * Reads up to {@code len} bytes into the given buffer, starting at {@code off}.
     *
     * @param b   the destination buffer
     * @param off the start offset
     * @param len the maximum number of bytes to read
     * @return the number of bytes read, or {@code -1} if end of stream
     * @throws IOException if an I/O or cipher processing error occurs
     */
    @Override
    public int read(final byte[] b, final int off, final int len) throws IOException {
        if (outputPos >= outputLen && !fillBuffer()) {
            return -1;
        }
        final int toCopy = Math.min(len, outputLen - outputPos);
        System.arraycopy(outputBuffer, outputPos, b, off, toCopy);
        outputPos += toCopy;
        return toCopy;
    }

    /**
     * Attempts to refill the output buffer by reading from the input stream and
     * processing the data with the cipher.
     * <p>
     * If the end of input is reached, this method finalizes the cipher (if
     * applicable) and marks the stream as complete.
     *
     * @return {@code true} if new output was generated and is available to read,
     *         {@code false} if end of stream was reached and no more output is
     *         available
     * @throws IOException if cipher processing fails
     */
    protected boolean fillBuffer() throws IOException {
        if (finalProcessed) {
            return false;
        }

        int read;
        try {
            do {
                read = in.readNBytes(inputBuffer, 0, inputBuffer.length);
                outputLen = processCipher(inputBuffer, read, outputBuffer);
            } while (outputLen == 0 && read > 0);
            finalProcessed = read == 0;
            outputPos = 0;
            return outputLen > 0;
        } catch (DataLengthException | IllegalStateException | InvalidCipherTextException e) {
            throw new IOException("Cipher processing failed", e);
        }
    }

    /**
     * Processes a chunk of data using the configured cipher.
     * <p>
     * Implementations must detect {@code inputLen == 0} to perform finalization if
     * the cipher supports it (e.g., for block or AEAD ciphers).
     * <p>
     * For stream ciphers, this typically means returning {@code 0} as no
     * finalization is needed.
     *
     * @param input    the buffer containing input data; contents are undefined if
     *                 {@code inputLen == 0}
     * @param inputLen the number of bytes to process, or {@code 0} to finalize the
     *                 cipher
     * @param output   the buffer into which processed output is written
     * @return the number of bytes written to the output buffer
     * 
     * @throws DataLengthException        if the input data is too large for the
     *                                    cipher
     * @throws IllegalStateException      if the cipher has not been properly
     *                                    initialized
     * @throws InvalidCipherTextException if finalization fails due to invalid
     *                                    padding or corrupted ciphertext
     * @throws Exception                  if any other unexpected error occurs
     *                                    during processing
     */
    protected abstract int processCipher(byte[] input, int inputLen, byte[] output) throws InvalidCipherTextException;

    /**
     * Closes the underlying input stream and releases any associated resources.
     *
     * @throws IOException if an I/O error occurs during closing
     */
    @Override
    public void close() throws IOException {
        in.close();
    }
}

/**
 * An {@link InputStream} that applies a {@link BufferedBlockCipher} to
 * transform the data during reading.
 * <p>
 * This stream supports both encryption and decryption depending on how the
 * cipher is initialized.
 */
class BufferedBlockCipherInputStream extends SymmetricCipherInputStreamBase {
    private final BufferedBlockCipher cipher;

    /**
     * Creates a new stream that wraps the given input and processes it using a
     * {@link BufferedBlockCipher}.
     *
     * @param in     the input stream to wrap
     * @param cipher the initialized {@link BufferedBlockCipher}
     */
    public BufferedBlockCipherInputStream(final InputStream in, final BufferedBlockCipher cipher) {
        super(in, cipher.getOutputSize(SymmetricStreamBuilder.BUF_SIZE));
        this.cipher = cipher;
    }

    /**
     * Processes a block of input using the {@link BufferedBlockCipher}.
     * <p>
     * If {@code inputLen == 0}, the cipher is finalized via
     * {@link BufferedBlockCipher#doFinal(byte[], int)}, completing any remaining
     * buffered input and writing final output (including padding, if applicable).
     * Otherwise, the method delegates to
     * {@link BufferedBlockCipher#processBytes(byte[], int, int, byte[], int)}.
     *
     * @param input    the input buffer containing plaintext or ciphertext bytes;
     *                 ignored when {@code inputLen == 0}
     * @param inputLen number of bytes to process from the input buffer, or
     *                 {@code 0} to trigger finalization
     * @param output   the output buffer to write processed data into
     * @return the number of bytes written to the output buffer
     *
     * @throws DataLengthException        if the input data is too large for the
     *                                    cipher
     * @throws IllegalStateException      if the cipher has not been properly
     *                                    initialized
     * @throws InvalidCipherTextException if finalization fails due to invalid
     *                                    padding or corrupted ciphertext
     */
    @Override
    protected int processCipher(final byte[] input, final int inputLen, final byte[] output)
            throws InvalidCipherTextException {
        if (inputLen == 0) {
            return cipher.doFinal(output, 0);
        }
        return cipher.processBytes(input, 0, inputLen, output, 0);
    }
}

/**
 * An {@link InputStream} that applies an {@link AEADBlockCipher} to transform
 * the data during reading.
 * <p>
 * AEAD (Authenticated Encryption with Associated Data) ciphers require
 * finalization to validate authentication tags.
 */
class AEADBlockCipherInputStream extends SymmetricCipherInputStreamBase {
    private final AEADBlockCipher cipher;

    /**
     * Creates a new stream that wraps the given input and processes it using an
     * {@link AEADBlockCipher}.
     *
     * @param in     the input stream to wrap
     * @param cipher the initialized {@link AEADBlockCipher}
     */
    public AEADBlockCipherInputStream(final InputStream in, final AEADBlockCipher cipher) {
        super(in, cipher.getOutputSize(SymmetricStreamBuilder.BUF_SIZE) + AesSupport.BLOCK_SIZE);
        this.cipher = cipher;
    }

    /**
     * Processes a block of input using the {@link AEADBlockCipher}.
     * <p>
     * If {@code inputLen == 0}, the cipher is finalized via
     * {@link AEADBlockCipher#doFinal(byte[], int)}, completing encryption or
     * decryption and verifying the authentication tag. If authentication fails, an
     * {@link InvalidCipherTextException} is thrown.
     * <p>
     * For normal processing, the method delegates to
     * {@link AEADBlockCipher#processBytes(byte[], int, int, byte[], int)} to
     * transform the input data into ciphertext or plaintext.
     *
     * @param input    the input buffer containing plaintext or ciphertext data;
     *                 ignored when {@code inputLen == 0}
     * @param inputLen number of bytes to process, or {@code 0} to finalize and
     *                 verify authentication
     * @param output   the output buffer to receive processed data
     * @return the number of bytes written to the output buffer
     *
     * @throws IllegalStateException      if the cipher is not properly initialized
     * @throws InvalidCipherTextException if finalization fails due to
     *                                    authentication tag mismatch
     */
    @Override
    protected int processCipher(final byte[] input, final int inputLen, final byte[] output)
            throws InvalidCipherTextException {
        if (inputLen == 0) {
            return cipher.doFinal(output, 0);
        }

        return cipher.processBytes(input, 0, inputLen, output, 0);
    }
}

/**
 * An {@link InputStream} that applies a {@link StreamCipher} to transform the
 * data during reading.
 * <p>
 * Stream ciphers operate byte-by-byte and do not require finalization.
 */
class StreamCipherInputStream extends SymmetricCipherInputStreamBase {
    private final StreamCipher cipher;

    /**
     * Creates a new stream that wraps the given input and processes it using a
     * {@link StreamCipher}.
     *
     * @param in     the input stream to wrap
     * @param cipher the initialized {@link StreamCipher}
     */
    public StreamCipherInputStream(final InputStream in, final StreamCipher cipher) {
        super(in, SymmetricStreamBuilder.BUF_SIZE);
        this.cipher = cipher;
    }

    /**
     * Processes a block of input using the stream cipher. Finalization is a no-op
     * for stream ciphers (returns 0 bytes).
     *
     * @param input    input buffer
     * @param inputLen number of bytes to process, or {@code 0} to signal end of
     *                 stream
     * @param output   output buffer
     * @return number of bytes written to output
     */
    @Override
    protected int processCipher(final byte[] input, final int inputLen, final byte[] output) {
        if (inputLen == 0) {
            // No finalization for StreamCipher, just indicate end of stream
            return 0;
        }
        cipher.processBytes(input, 0, inputLen, output, 0);
        return inputLen;
    }
}