#version 430 core

layout (local_size_x = 6, local_size_y = 6) in;

layout (binding = 0, rgba8) writeonly uniform image2D destImg;

layout (binding = 0) buffer AstcSource {
    uint astc[];
};

const uvec2 block_size = gl_WorkGroupSize.xy;

const uint block_width  = block_size.x;
const uint block_height = block_size.y;

void assert(bool a) {}

uint bit_pop_cnt(uint number) {
    uint counter;
    for (counter = 0; number != 0; counter++) {
        number &= number - 1;
    }
    return counter;
}

bool read_bit(uint stream) {
    uint word = astc[stream / 32];
    return (word & (stream % 32)) != 0;
}

uint read_astc_bits(inout uint stream, uint len) {
    uint position = stream;
    uint ret_value = 0;
    for (uint i = position; i < position + len; i++) {
        if (read_bit(i)) {
            ret_value |= 1 << (i - position);
        }
    }
    stream += len;
    return ret_value;
}

void write_bits(uint size, uint value, inout uint bytes[4], inout uint cursor) {
    uint position = cursor;
    for (uint i = position; i < position + size; i++) {
        if ((value & (i - position)) != 0) {
            bytes[i / 32] |= 1 << (i % 32);
        }
    }
    cursor += size;
}

const uint INTEGER_ENCODING_JUST_BITS = 0u;
const uint INTEGER_ENCODING_QUINT     = 1u;
const uint INTEGER_ENCODING_TRIT      = 2u;

struct IntegerEncoding {
    uint encoding;
    uint number_bits;
    uint bit_value;
    uint trit_value;
    uint quint_value;
};

IntegerEncoding create_integer_encoding(uint max_val) {
    IntegerEncoding encoding;
    encoding.bit_value = 0;
    encoding.trit_value = 0;
    encoding.quint_value = 0;

    while (max_val > 0) {
        uint check = max_val + 1;

        //is max_val a power of two?
        if ((check & (check - 1)) == 0) {
            encoding.encoding = INTEGER_ENCODING_JUST_BITS;
            encoding.number_bits = bit_pop_cnt(max_val);
            return encoding;
        //is max_val of the type 3*2^n - 1?
        } else if ((check % 3 == 0) && ((check / 3) & ((check / 3) - 1)) == 0) {
            encoding.encoding = INTEGER_ENCODING_TRIT;
            encoding.number_bits = bit_pop_cnt(check / 3 - 1);
            return encoding;
        //is max_val of the type 5*2^n - 1?
        } else if ((check % 5 == 0) && (((check / 5) & (check / 5) - 1)) == 0) {
            encoding.encoding = INTEGER_ENCODING_QUINT;
            encoding.number_bits = bit_pop_cnt(check / 5 - 1);
            return encoding;
        }
        //apparently it can't be represented with a bounded integer sequence; just iterate
        max_val--;
    }
    encoding.encoding = INTEGER_ENCODING_JUST_BITS;
    encoding.number_bits = 0;
    return encoding;
}

uint get_bit_length(IntegerEncoding encoding, uint number_vals) {
    uint total_bits = encoding.number_bits * number_vals;
    if (encoding.encoding == INTEGER_ENCODING_TRIT) {
        total_bits += (number_vals * 8 + 4) / 5;
    } else if (encoding.encoding == INTEGER_ENCODING_QUINT) {
        total_bits += (number_vals * 7 + 2) / 3;
    }
    return total_bits;
}

struct TexelWeightParams {
    uint width;
    uint height;
    bool dual_plane;
    uint max_weight;
    bool error;
    bool void_extent_ldr;
    bool void_extent_hdr;
};

uint get_packed_bit_size(TexelWeightParams params) {
    //how many indices do we have?
    uint indices = params.width * params.height;
    if (params.dual_plane) {
        indices *= 2;
    }
    IntegerEncoding encoding = create_integer_encoding(params.max_weight);
    return get_bit_length(encoding, indices);
}

TexelWeightParams decode_block_info(inout uint stream) {
    TexelWeightParams params;

    //read the entire block mode all at once
    uint mode_bits = read_astc_bits(stream, 11);

    //does this match the void extent block mode?
    if ((mode_bits & 0x01ff) == 0x1fc) {
        if ((mode_bits & 0x200) != 0) {
            params.void_extent_hdr = true;
        } else {
            params.void_extent_ldr = true;
        }
        //next two bits must be one
        if (((mode_bits & 0x400) == 0) || (read_astc_bits(stream, 1) == 0)) {
            params.error = true;
        }
        return params;
    }

    //first check if the last four bits are zero
    if ((mode_bits & 0xf) == 0) {
        params.error = true;
        return params;
    }

    //if the last two bits are zero, then if bits [6-8] are all ones, this is also reserved
    if ((mode_bits & 0x3) == 0 && (mode_bits & 0x1C0) == 0x1C0) {
        params.error = true;
        return params;
    }

    //otherwise, there is no error... figure out the layout of the block mode
    //layout is determined by a number between 0 and 9 corresponding to table C.2.8 of the ASTC spec
    uint layaout = 0;

    if ((mode_bits & 0x1) != 0 || (mode_bits & 0x2) != 0) {
        //layout is in [0-4]
        if ((mode_bits & 0x8) != 0) {
            //layout is in [2-4]
            if ((mode_bits & 0x4) != 0) {
                // layout is in [3-4]
                if ((mode_bits & 0x100) != 0) {
                    layaout = 4;
                } else {
                    layaout = 3;
                }
            } else {
                layaout = 2;
            }
        } else {
            //layout is in [0-1]
            if ((mode_bits & 0x4) != 0) {
                layaout = 1;
            } else {
                layaout = 0;
            }
        }
    } else {
        if ((mode_bits & 0x100) != 0) {
            //layout is in [7-9]
            if ((mode_bits & 0x80) != 0) {
                // layout is in [7-8]
                assert((mode_bits & 0x40) == 0);

                if ((mode_bits & 0x20) != 0) {
                    layaout = 8;
                } else {
                    layaout = 7;
                }
            } else {
                layaout = 9;
            }
        } else {
            // layout is in [5-6]
            if ((mode_bits & 0x80) != 0) {
                layaout = 6;
            } else {
                layaout = 5;
            }
        }
    }

    assert(layaout < 10);

    //determine r
    uint r = (mode_bits >> 4) & 1;
    if (layaout < 5) {
        r |= (mode_bits & 0x3) << 1;
    } else {
        r |= (mode_bits & 0xc) >> 1;
    }

    assert(2 <= r && r <= 7);

    uint a, b;
    switch (layaout) {
        case 0:
            a = (mode_bits >> 5) & 0x3;
            b = (mode_bits >> 7) & 0x3;
            params.width  = b + 4;
            params.height = a + 2;
            break;
        case 1:
            a = (mode_bits >> 5) & 0x3;
            b = (mode_bits >> 7) & 0x3;
            params.width  = b + 8;
            params.height = a + 2;
            break;
        case 2:
            a = (mode_bits >> 5) & 0x3;
            b = (mode_bits >> 7) & 0x3;
            params.width  = a + 2;
            params.height = b + 8;
            break;
        case 3:
            a = (mode_bits >> 5) & 0x3;
            b = (mode_bits >> 7) & 0x1;
            params.width  = a + 2;
            params.height = b + 6;
            break;
        case 4:
            a = (mode_bits >> 5) & 0x3;
            b = (mode_bits >> 7) & 0x1;
            params.width  = b + 2;
            params.height = a + 2;
            break;
        case 5:
            a = (mode_bits >> 5) & 0x3;
            params.width  = 12;
            params.height = a + 2;
            break;
        case 6:
            a = (mode_bits >> 5) & 0x3;
            params.width  = a + 2;
            params.height = 12;
            break;
        case 7:
            params.width  = 6;
            params.height = 10;
            break;
        case 8:
            params.width  = 10;
            params.height = 6;
            break;
        case 9:
            a = (mode_bits >> 5) & 0x3;
            b = (mode_bits >> 9) & 0x3;
            params.width  = a + 6;
            params.height = b + 6;
            break;
        default:
            params.error = true;
            break;
    }
    bool d = ((layaout != 9) && (mode_bits & 0x400) != 0);
    bool h = (layaout != 9) && ((mode_bits & 0x200) != 0);

    if (h) {
        const int max_weights[] = { 9, 11, 15, 19, 23, 31 };
        params.max_weight = max_weights[r - 2];
    } else {
        const int max_weights[] = { 1, 2, 3, 4, 5, 7 };
        params.max_weight = max_weights[r - 2];
    }
    params.dual_plane = d;

    return params;
}

bool decompress_block(uint input_offset) {
    uint stream = input_offset * 8;
    TexelWeightParams params = decode_block_info(stream);

    assert(!params.error);

    if (params.void_extent_ldr) {
        fill_void_extent_ldr(stream);
        return true;
    }

    assert(!params.void_extent_hdr);
    assert(params.width > block_width);
    assert(params.height > block_height);

    //read num partitions
    uint number_partitions = read_astc_bits(stream, 2) + 1;
    assert(number_partitions <= 4);

    assert(!(number_partitions == 4 && params.dual_plane));

    uint partition_index;
    uint color_endpoint_mode[] = { 0, 0, 0, 0 };

    //read extra config data
    uint base_color_endpoint_mode = 0;

    if (number_partitions == 1) {
        color_endpoint_mode[0] = read_astc_bits(stream, 4);
        partition_index = 0;
    } else {
        partition_index = read_astc_bits(stream, 10);
        base_color_endpoint_mode = read_astc_bits(stream, 6);
    }

    uint base_mode = base_color_endpoint_mode & 3;

    //remaining bits are color endpoint data...
    uint number_weight_bits = get_packed_bit_size(params);
    uint remaining_bits = 128 - number_weight_bits - stream;

    uint extra_color_endpoint_mode_bits = 0;

    if (base_mode != 0) {
        switch (number_partitions) {
            case 2:  extra_color_endpoint_mode_bits += 2; break;
            case 3:  extra_color_endpoint_mode_bits += 5; break;
            case 4:  extra_color_endpoint_mode_bits += 8; break;
            default: assert(false); break;
        }
    }

    remaining_bits -= extra_color_endpoint_mode_bits;

    //do we have a dual plane situation?
    uint plane_selector_bits = 0;
    if (params.dual_plane) {
        plane_selector_bits = 2;
    }
    remaining_bits -= plane_selector_bits;

    //read color data
    uint color_data_bits = remaining_bits;

    //color endpoint
    uint cep[4];
    uint cep_cursor = 0;

    while (remaining_bits > 0) {
        uint number_bits = min(remaining_bits, 8);
        uint bits = read_astc_bits(stream, number_bits);
        write_bits(bits, number_bits, cep, cep_cursor);
        remaining_bits -= 8;
    }

    //read the plane selection bits
    uint plane_indices = read_astc_bits(stream, plane_selector_bits);

    //read the rest of the CEM
    if (base_mode != 0) {
        uint extra_color_endpoint_mode = read_astc_bits(stream, extra_color_endpoint_mode_bits);
        uint temp_color_endpoint_mode = (extra_color_endpoint_mode << 6) | base_color_endpoint_mode;
        temp_color_endpoint_mode >>= 2;

        bool c[] = { false, false, false, false };
        for (uint i = 0; i < number_partitions; i++) {
            c[i] = (temp_color_endpoint_mode & 1) != 0;
            temp_color_endpoint_mode >>= 1;
        }

        uint m[] = { 0, 0, 0, 0 };
        for (uint i = 0; i < number_partitions; i++) {
            m[i] = uint(temp_color_endpoint_mode & 3);
            temp_color_endpoint_mode >>= 2;
            assert(m[i] <= 3);
        }

        for (uint i = 0; i < number_partitions; i++) {
            color_endpoint_mode[i] = base_mode;
            if (!(c[i])) color_endpoint_mode[i]--;
            color_endpoint_mode[i] <<= 2;
            color_endpoint_mode[i] |= m[i];
        }
    } else if (number_partitions > 1) {
        uint temp_color_endpoint_mode = base_color_endpoint_mode >> 2;
        for (uint i = 0; i < number_partitions; i++) {
            color_endpoint_mode[i] = temp_color_endpoint_mode;
        }
    }

    //make sure everything is up till here is name
    for (uint i = 0; i < number_partitions; i++) {
        assert(color_endpoint_mode[i] < 16);
    }
    assert(stream + get_packed_bit_size(params) == 128);

    //continue in line 278 ASTCDecoder.cs
    //...
}

void decode() {
    uint j = gl_WorkGroupID.x * block_size.x;
    uint i = gl_WorkGroupID.y * block_size.y;


}

void main() {
    uvec2 pixel_coord = gl_WorkGroupID.xy * gl_WorkGroupSize.xy + gl_LocalInvocationID.xy;

    vec2 colors = vec2(gl_LocalInvocationID.xy) / vec2(gl_WorkGroupSize.xy);

    vec4 final_color = vec4(colors, 0.0, 1.0);

    imageStore(destImg, ivec2(pixel_coord), final_color);
}