#include <PSX/Auxiliary/lz4_decompressor.hpp>
#include <PSX/Auxiliary/unaligned_read.hpp>

namespace JabyEngine {
    static void memcpy(uint8_t* &dst, ArrayRange<const uint8_t> &src, size_t size) {
        for(size_t n = 0; n < size; n++) {
            *dst = src.pop();
            dst++;
        }
    }

    bool LZ4Decompressor :: obtain_any_length(ArrayRange<const uint8_t>& data, uint32_t &dst_length) {
        while(data) {
            const auto additional_value = data.pop();
            
            dst_length += additional_value;
            if(additional_value != 255) {
                return true;
            }
        }

        return false;
    }

    void LZ4Decompressor :: setup(uint8_t* dst_adr) {
        this->dst_adr = dst_adr;
        this->state.enable();
    }

    void LZ4Decompressor :: disable() {
        this->state.disable();
    }

    LZ4Decompressor::Result LZ4Decompressor :: process(ArrayRange<const uint8_t> data, bool is_last) {
        size_t bytes_ready = 0;

        while(data) {
            switch(this->state.step) {
                case State::Step::Disabled:
                    return Result::new_done(data.size);

                case State::Step::ReadToken: {
                    const auto token = data.pop();

                    this->state.literal_length = (token & 0xF0) >> 4;
                    this->state.match_length   = (token & 0x0F);

                    if(this->state.literal_length == 15) {
                        this->state.step = State::Step::ObtainLiteralLength;
                    }

                    else if(this->state.literal_length == 0) {
                        this->state.step = State::Step::ObtainMatchOffset;
                    }

                    else {
                        this->state.step = State::Step::CopyLiterals;
                    }
                } break;

                case State::Step::ObtainLiteralLength: {
                    if(LZ4Decompressor::obtain_any_length(data, this->state.literal_length)) {
                        this->state.step = State::Step::CopyLiterals;            
                    }
                } break;

                case State::Step::CopyLiterals: {
                    const auto bytes_copy = (this->state.literal_length > data.size) ? data.size : this->state.literal_length;

                    memcpy(this->dst_adr, data, bytes_copy);

                    this->state.literal_length -= bytes_copy;
                    bytes_ready += bytes_copy;
                    
                    if(this->state.literal_length == 0) {
                        this->state.step = State::Step::ObtainMatchOffset;    
                    }
                } break;

                case State::Step::ObtainMatchOffset: {
                    static const auto state_complete = [](State& state) -> bool {
                        if(state.match_offset == 0) {
                            return false;
                        }

                        else {
                            if(state.match_length == 15) {
                                state.step = State::Step::ObtainMatchLength;
                            }

                            else {
                                state.step = State::Step::CopyMatch;
                            }
                            return true;
                        }
                    };

                    if(this->state.match_offset == 0xFFFF) {
                        // We are unused and invalid
                        if(data.size >= sizeof(uint16_t)) {
                            // We can read all
                            this->state.match_offset = unaligned_lhu(data.start);
                            data.skip(sizeof(uint16_t));

                            if(!state_complete(this->state)) {
                                return Result::new_done(bytes_ready);
                            }
                        }

                        else {
                            this->state.match_offset = static_cast<uint16_t>(data.pop());
                        }
                    }

                    else {
                        this->state.match_offset |= (static_cast<uint16_t>(data.pop()) << 8);
                        if(!state_complete(this->state)) {
                            return Result::new_done(bytes_ready);
                        }
                    }
                } break;

                case State::Step::ObtainMatchLength: {
                    if(LZ4Decompressor::obtain_any_length(data, this->state.match_length)) {
                        this->state.step = State::Step::CopyMatch;
                    }
                } break;

                case State::Step::CopyMatch: {
                    static constexpr size_t min_match_length = 4;

                    this->state.match_length += min_match_length;

                    const uint8_t* src = this->dst_adr - this->state.match_offset;
                    ArrayRange     src_data(src, this->state.match_length);

                    memcpy(this->dst_adr, src_data, this->state.match_length);

                    this->state.match_offset = 0xFFFF;
                    this->state.step         = State::Step::ReadToken;

                    bytes_ready += this->state.match_length;
                } break;

                default:
                    return Result::new_error();
            }
        }
        return is_last ? Result::new_done(bytes_ready) : Result::new_in_progress(bytes_ready);
    }
}