package cn.skcks.matrix.v2.utils.base91;

import java.io.ByteArrayOutputStream;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;

public class Base91 {
    private static final float AVERAGE_ENCODING_RATIO = 1.2297f;
    private static final int BASE = 91;
    private static final char[] ENCODING_TABLE = {
            'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N',
            'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
            'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
            'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
            '0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
            '!', '#', '$', '%', '&', '(', ')', '*', '+', ',', '.', '/', ':', ';',
            '<', '=', '>', '?', '@', '[', ']', '^', '_', '`', '{', '|', '}', '~',
            '"'};

    private static final int[] DECODING_TABLE;

    static {
        int[] charts = new int[255];
        Arrays.fill(charts, -1);
        for (int i = 0; i < ENCODING_TABLE.length; i++) {
            charts[ENCODING_TABLE[i]] = i;
        }
        DECODING_TABLE = charts;
    }

    public static String encodeString(String str, Charset charset) {
        return encode(str.getBytes(charset));
    }

    public static String encodeString(String str) {
        return encode(str.getBytes(StandardCharsets.UTF_8));
    }

    public static String encode(byte[] bytes) {
        int ebq = 0;
        int en = 0;
        StringBuilder buffer = new StringBuilder();
        for (int b : bytes) {
            ebq |= (b & 255) << en;
            en += 8;
            if (en > 13) {
                int ev = ebq & 8191;

                if (ev > 88) {
                    ebq >>= 13;
                    en -= 13;
                } else {
                    ev = ebq & 16383;
                    ebq >>= 14;
                    en -= 14;
                }
                buffer.append(ENCODING_TABLE[ev % BASE]);
                buffer.append(ENCODING_TABLE[ev / BASE]);
            }
        }

        if (en > 0) {
            buffer.append(ENCODING_TABLE[ebq % BASE]);
            if (en > 7 || ebq > 90) {
                buffer.append(ENCODING_TABLE[ebq / BASE]);
            }
        }

        return buffer.toString();
    }

    public static byte[] decode(String str) {
        byte[] data = str.getBytes(StandardCharsets.UTF_8);
        int dbq = 0, dn = 0, dv = -1;
        int estimatedSize = Math.round(data.length / AVERAGE_ENCODING_RATIO);
        ByteArrayOutputStream buffer = new ByteArrayOutputStream(estimatedSize);
        for (byte datum : data) {
            if (dv == -1)
                dv = DECODING_TABLE[datum];
            else {
                dv += DECODING_TABLE[datum] * BASE;
                dbq |= dv << dn;
                dn += (dv & 8191) > 88 ? 13 : 14;
                do {
                    buffer.write(dbq);
                    dbq >>= 8;
                    dn -= 8;
                } while (dn > 7);
                dv = -1;
            }
        }
        if (dv != -1) {
            buffer.write((dbq | dv << dn));
        }
        return buffer.toByteArray();
    }

    public static String decodeString(String str) {
        return new String(decode(str), StandardCharsets.UTF_8);
    }

    public static String decodeString(String str, Charset charset) {
        return new String(decode(str), charset);
    }
}
