package com.atlassian.bitbucket.mesh.io;

import org.mozilla.universalchardet.UniversalDetector;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.nio.charset.Charset;

import static com.google.common.base.Preconditions.checkArgument;
import static java.util.Objects.requireNonNull;

public class EncodingUtils {

    /**
     * Tries to detect the encoding of the buffered data, first using byte order marks if present and then by
     * applying heuristics to the bytes. The provided buffer must contain at least 4 bytes or detection will
     * not even be attempted.
     *
     * @param buffer the buffer to detect an encoding for
     * @return the detected character set, which may be {@code null} if an encoding could not be determined
     * @throws NullPointerException if the provided {@code buffer} is {@code null}
     */
    @Nullable
    public static String detectEncoding(@Nonnull ByteBuffer buffer) {
        if (requireNonNull(buffer, "buffer").remaining() < 4) {
            return null;
        }

        byte[] buf = new byte[Math.min(buffer.remaining(), 4)];
        buffer.mark();
        buffer.get(buf);
        buffer.reset();

        UniversalDetector detector = new UniversalDetector(null);
        if (buf.length > 3) { //JUniversalChardet's BOM detection requires at least 4 bytes
            detector.handleData(buf, 0, buf.length);
            if (detector.isDone()) {
                //If a BOM is present, the first 4 bytes will be enough to detect the charset
                return detector.getDetectedCharset();
            }

            //If we make it here, that means the file doesn't begin with byte order marks specifying
            //UTF-8, 16 or 32. That means we need to detect a bit harder.
            detector.reset(); //Reset the detector; it's going to get the same 4 first bytes over
        }

        if (buffer.hasArray()) {
            //If the buffer has an accessible byte array (non-read-only HeapByteBuffer), reference the
            //array directly to avoid the overhead of having to copy it
            detector.handleData(buffer.array(), buffer.position(), buffer.remaining());
        } else {
            //If the buffer doesn't have an accessible byte array (DirectByteBuffer, or a read-only
            //HeapByteBuffer), copy the buffer into a new byte[] and feed that to the detector
            buf = copyBuffer(buffer);

            detector.handleData(buf, 0, buf.length);
        }
        detector.dataEnd();

        return getDetectedCharset(detector);
    }

    /**
     * Tries to detect the encoding of the data, first using byte order marks if present and then by applying
     * heuristics to the first {@code bufferSize} bytes. If the specified {@code bufferSize} is less than 4
     * bytes, detection will not be attempted.
     *
     * @param stream     the stream containing the data, which must support {@code mark(int)} and {@code reset()}
     * @param bufferSize the number of bytes to read from the stream when applying heuristic type detection
     * @return the detected character set, which may be {@code null} if the detector couldn't determine what encoding was used
     * @throws IOException if bytes cannot be read from the provided {@code stream}
     * @throws NullPointerException if the provided {@code stream} is {@code null}
     */
    @Nullable
    public static String detectEncoding(@Nonnull InputStream stream, int bufferSize) throws IOException {
        requireNonNull(stream, "stream");
        checkArgument(stream.markSupported(), "Encoding detection requires a stream which supports mark/reset.");
        checkArgument(bufferSize > 0, "Encoding detection requires a positive buffer size.");
        if (bufferSize < 4) {
            return null;
        }

        //The first check, we only read in the first 4 bytes to check for the presence of byte order marks.
        byte[] buf = new byte[4];
        stream.mark(5);
        int read = stream.read(buf);
        stream.reset();

        UniversalDetector detector = new UniversalDetector(null);
        if (read > 3) { //JUniversalChardet's BOM detection requires at least 4 bytes
            detector.handleData(buf, 0, read);
            if (detector.isDone()) {
                return detector.getDetectedCharset();
            }

            //If we make it here, that means the file doesn't begin with byte order marks specifying a UTF-8, 16 or
            //32. That means we need to detect a bit harder.
            detector.reset(); //Reset the detector; it's going to get the same 4 first bytes over
        }

        buf = new byte[bufferSize];
        stream.mark(buf.length); //Mark the stream so we can reset for reading content

        //Spool over the file until we've read out the maximum number of bytes we're prepared to, or until the
        //detection logic decides it's figured out the character set
        int remaining = buf.length;
        while (remaining > 0 && !detector.isDone() && (read = stream.read(buf, 0, remaining)) > 0) {
            detector.handleData(buf, 0, read);
            remaining -= read;
        }
        stream.reset(); //Reset the stream so the caller can read content
        detector.dataEnd(); //We're not going to provide any more bytes, so go with the best guess

        //Return whatever we decided (which may be null)
        return getDetectedCharset(detector);
    }

    /**
     * Tries to detect the encoding of the data by applying heuristics over the entire byte array. The provided array
     * must contain at least 4 bytes or detection will not even be attempted.
     *
     * @param buffer the bytes containing the data
     * @return the detected character set, which may be {@code null} if the detector couldn't determine what encoding was used
     */
    @Nullable
    public static String detectEncoding(@Nonnull byte[] buffer) {
        if (requireNonNull(buffer, "buffer").length < 4) {
            return null;
        }

        UniversalDetector detector = new UniversalDetector(null);
        detector.handleData(buffer, 0, buffer.length);
        detector.dataEnd();

        return getDetectedCharset(detector);
    }

    /**
     * Tries to detect whether the data is binary or not. When there is no clear indication of whether the content is
     * binary or text, {@code false} is returned.
     *
     * @param buffer a buffer containing bytes to check
     * @return {@code true} if binary is detected, {@code false} otherwise
     */
    public static boolean isBinary(@Nonnull ByteBuffer buffer) {
        if (!requireNonNull(buffer, "buffer").hasRemaining()) {
            return false;
        }

        if (buffer.hasArray()) {
            return isBinary(buffer.array(), buffer.position() + buffer.arrayOffset(), buffer.remaining());
        }

        byte[] buf = copyBuffer(buffer);
        return isBinary(buf, 0, buf.length);
    }

    /**
     * Tries to detect whether the data is binary or not. When there is no clear indication of whether the content is
     * binary or text, {@code false} is returned.
     *
     * @param stream     the stream containing the data, which must support {@code mark(int)} and {@code reset()}
     * @param bufferSize the maximum number of bytes to read from the stream for detection (must be greater than 0)
     * @return {@code true} if binary is detected, {@code false} otherwise
     * @throws IOException if bytes cannot be read from the provided {@code stream}
     */
    public static boolean isBinary(@Nonnull InputStream stream, int bufferSize) throws IOException {
        requireNonNull(stream, "stream");
        checkArgument(stream.markSupported(), "Binary detection requires a stream which supports mark/reset.");
        checkArgument(bufferSize > 0, "Binary detection requires a positive buffer size.");

        byte[] buf = new byte[4];
        stream.mark(5);
        int read = stream.read(buf);
        stream.reset();

        if (read > 3) { //JUniversalChardet's BOM detection requires at least 4 bytes
            UniversalDetector detector = new UniversalDetector(null);
            detector.handleData(buf, 0, read);
            if (detector.isDone()) {
                //If a charset could be detected from the provided bytes, assume it's not binary
                return false;
            }
        }

        //No BOM; check for NULs
        stream.mark(bufferSize);
        read = 0;
        int b;
        while (read < bufferSize && (b = stream.read()) != -1) {
            if (b == 0) {
                stream.reset();
                return true;
            }
            ++read;
        }
        stream.reset();
        return false;
    }

    /**
     * Tries to detect whether the data is binary or not. When there is no clear indication of whether the content is
     * binary or text, {@code false} is returned.
     *
     * @param buffer the bytes containing the data
     * @return {@code true} if binary is detected, {@code false} otherwise
     */
    public static boolean isBinary(@Nonnull byte[] buffer) {
        return isBinary(buffer, 0, buffer.length);
    }

    /**
     * Marks the provided buffer and then copies all of its bytes to an array using a single {@code get(byte[])}.
     * For direct {@link ByteBuffer}s, this is substantially more performant than calling {@code get()} for each
     * byte, since each call requires a roundtrip to native memory.
     *
     * @param buffer the buffer to copy, which will be marked and reset
     * @return the copied bytes
     */
    private static byte[] copyBuffer(@Nonnull ByteBuffer buffer) {
        buffer.mark();

        try {
            byte[] buf = new byte[buffer.remaining()];
            buffer.get(buf);

            return buf;
        } finally {
            buffer.reset();
        }
    }

    private static String getDetectedCharset(UniversalDetector detector) {
        String encoding = detector.getDetectedCharset();
        //JUniversalChardet returns "US-ASCII" if no other charset was detected and the data is "printable"
        //(contains only the standard ASCII characters plus new line, carriage return and tab).
        //For our purposes this is treated the same as if none was detected, so return null in that case.
        if (encoding != null && Charset.isSupported(encoding) && !"US-ASCII".equals(encoding)) {
            //If a charset is detected, it's only returned if the JVM supports it. Otherwise, it's better to
            //let callers default (most likely to UTF-8) and hope for the best
            return encoding;
        }

        return null;
    }

    private static boolean isBinary(byte[] buffer, int offset, int count) {
        requireNonNull(buffer, "buffer");

        if (count > 3) { //JUniversalChardet's BOM detection requires at least 4 bytes
            UniversalDetector detector = new UniversalDetector(null);
            detector.handleData(buffer, offset, 4);
            if (detector.isDone()) {
                //If a charset could be detected from the provided bytes, assume it's not binary
                return false;
            }
        }

        //No BOM; check for NULs
        int end = offset + count;
        for (int i = offset; i < end; ++i) {
            if (buffer[i] == 0) {
                return true;
            }
        }
        return false;
    }
}