#ifndef _QS2_QD_DESERIALIZER_H_
#define _QS2_QD_DESERIALIZER_H_


#include <Rcpp.h>
#include <R_ext/Utils.h>

#include <Rversion.h>

#include "qx_file_headers.h"
#include "io/io_common.h"

using namespace Rcpp;


#if R_VERSION >= R_Version(4, 6, 0)

SEXP get_first_attrib(SEXP object) {
    return R_mapAttrib(object, [](SEXP, SEXP a, void*) {return a;}, nullptr);
}

struct DelayedAttribAssign {
    std::vector<SEXP> objects;
    void add(SEXP object) {
        objects.push_back(object);
    }
    void resolve() {
        if(objects.size() > 0) {
            PROTECT_INDEX prot_idx;
            for(size_t i=0; i < objects.size(); ++i) {
                SEXP object = objects[i];
                SEXP aobj = get_first_attrib(object);
                if(i == 0) {
                    PROTECT_WITH_INDEX(aobj, &prot_idx);
                } else {
                    REPROTECT(aobj, prot_idx);
                }
                Rf_setAttrib(object, R_SpecSymbol, R_NilValue); // remove placeholder aobj from object
                for (SEXP p = aobj; p != R_NilValue; p = CDR(p)) {
                    SEXP tag = TAG(p);
                    SEXP val = CAR(p);
                    Rf_setAttrib(object, tag, val);
                }
            }
            UNPROTECT(1);
        }
    }
};
#endif

template<typename block_compress_reader>
struct QdataDeserializer {
    block_compress_reader & reader;
    std::vector<std::pair<SEXP, uint64_t>> character_sexp;
    std::vector<std::pair<SEXP, uint64_t>> complex_sexp;
    std::vector<std::pair<SEXP, uint64_t>> real_sexp;
    std::vector<std::pair<SEXP, uint64_t>> integer_sexp; // and logical
    std::vector<std::pair<SEXP, uint64_t>> raw_sexp;

#if R_VERSION >= R_Version(4, 6, 0)
    DelayedAttribAssign delayed_attributes;
#endif

    QdataDeserializer(block_compress_reader & reader) : reader(reader) {}

    private:
    static constexpr uint64_t max_r_vector_length = static_cast<uint64_t>(R_XLEN_T_MAX);
    static constexpr uint32_t max_r_pairlist_length = static_cast<uint32_t>(R_LEN_T_MAX);
    static constexpr uint32_t max_r_string_length = static_cast<uint32_t>(R_LEN_T_MAX);

#ifndef LONG_VECTOR_SUPPORT
    void validate_object_length_32(const uint32_t length, const char * const what) {
        if(static_cast<uint64_t>(length) > max_r_vector_length) {
            reader.cleanup_and_throw(std::string(what) + " exceeds R_xlen_t range");
        }
    }
#endif

    void validate_object_length_64(const uint64_t length, const char * const what) {
        if(length > max_r_vector_length) {
            reader.cleanup_and_throw(std::string(what) + " exceeds R_xlen_t range");
        }
    }

    void validate_attr_length_32(const uint32_t length) {
        if(length > max_r_pairlist_length) {
            reader.cleanup_and_throw("Attribute count exceeds R pairlist length range");
        }
    }

    void validate_string_length_32(const uint32_t length, const char * const what) {
        if(length > max_r_string_length) {
            reader.cleanup_and_throw(std::string(what) + " exceeds R string size limit");
        }
    }

    // will also call get_pod_contiguous for object lengths
    void read_header_impl(const uint8_t header_byte, qstype & type, uint64_t & len) {
        uint8_t header_byte_5 = header_byte & bitmask_type_5;
        if(header_byte_5 == 0) {
            switch(header_byte) {
                case nil_header:
                    type = qstype::NIL;
                    return;
                case logical_header_8:
                    type = qstype::LOGICAL;
                    len = reader.template get_pod_contiguous<uint8_t>();
                return;
                case logical_header_16:
                    type = qstype::LOGICAL;
                    len = reader.template get_pod_contiguous<uint16_t>();
                    return;
                case logical_header_32:
                {
                    type = qstype::LOGICAL;
                    const uint32_t length = reader.template get_pod_contiguous<uint32_t>();
#ifndef LONG_VECTOR_SUPPORT
                    validate_object_length_32(length, "Logical vector length");
#endif
                    len = length;
                    return;
                }
                case logical_header_64:
                    type = qstype::LOGICAL;
                    len = reader.template get_pod_contiguous<uint64_t>();
                    validate_object_length_64(len, "Logical vector length");
                return;
                case integer_header_8:
                    type = qstype::INTEGER;
                    len = reader.template get_pod_contiguous<uint8_t>();
                    return;
                case integer_header_16:
                    type = qstype::INTEGER;
                    len = reader.template get_pod_contiguous<uint16_t>();
                    return;
                case integer_header_32:
                {
                    type = qstype::INTEGER;
                    const uint32_t length = reader.template get_pod_contiguous<uint32_t>();
#ifndef LONG_VECTOR_SUPPORT
                    validate_object_length_32(length, "Integer vector length");
#endif
                    len = length;
                    return;
                }
                case integer_header_64:
                    type = qstype::INTEGER;
                    len = reader.template get_pod_contiguous<uint64_t>();
                    validate_object_length_64(len, "Integer vector length");
                    return;
                case numeric_header_8:
                    type = qstype::REAL;
                    len = reader.template get_pod_contiguous<uint8_t>();
                    return;
                case numeric_header_16:
                    type = qstype::REAL;
                    len = reader.template get_pod_contiguous<uint16_t>();
                    return;
                case numeric_header_32:
                {
                    type = qstype::REAL;
                    const uint32_t length = reader.template get_pod_contiguous<uint32_t>();
#ifndef LONG_VECTOR_SUPPORT
                    validate_object_length_32(length, "Real vector length");
#endif
                    len = length;
                    return;
                }
                case numeric_header_64:
                    type = qstype::REAL;
                    len = reader.template get_pod_contiguous<uint64_t>();
                    validate_object_length_64(len, "Real vector length");
                    return;
                case complex_header_32:
                {
                    type = qstype::COMPLEX;
                    const uint32_t length = reader.template get_pod_contiguous<uint32_t>();
#ifndef LONG_VECTOR_SUPPORT
                    validate_object_length_32(length, "Complex vector length");
#endif
                    len = length;
                    return;
                }
                case complex_header_64:
                    type = qstype::COMPLEX;
                    len = reader.template get_pod_contiguous<uint64_t>();
                    validate_object_length_64(len, "Complex vector length");
                    return;
                case character_header_8:
                    type = qstype::CHARACTER;
                    len = reader.template get_pod_contiguous<uint8_t>();
                    return;
                case character_header_16:
                    type = qstype::CHARACTER;
                    len = reader.template get_pod_contiguous<uint16_t>();
                    return;
                case character_header_32:
                {
                    type = qstype::CHARACTER;
                    const uint32_t length = reader.template get_pod_contiguous<uint32_t>();
#ifndef LONG_VECTOR_SUPPORT
                    validate_object_length_32(length, "Character vector length");
#endif
                    len = length;
                    return;
                }
                case character_header_64:
                    type = qstype::CHARACTER;
                    len = reader.template get_pod_contiguous<uint64_t>();
                    validate_object_length_64(len, "Character vector length");
                    return;
                case list_header_8:
                    type = qstype::LIST;
                    len = reader.template get_pod_contiguous<uint8_t>();
                    return;
                case list_header_16:
                    type = qstype::LIST;
                    len = reader.template get_pod_contiguous<uint16_t>();
                    return;
                case list_header_32:
                {
                    type = qstype::LIST;
                    const uint32_t length = reader.template get_pod_contiguous<uint32_t>();
#ifndef LONG_VECTOR_SUPPORT
                    validate_object_length_32(length, "List length");
#endif
                    len = length;
                    return;
                }
                case list_header_64:
                    type = qstype::LIST;
                    len = reader.template get_pod_contiguous<uint64_t>();
                    validate_object_length_64(len, "List length");
                    return;
                case raw_header_32:
                {
                    type = qstype::RAW;
                    const uint32_t length = reader.template get_pod_contiguous<uint32_t>();
#ifndef LONG_VECTOR_SUPPORT
                    validate_object_length_32(length, "Raw vector length");
#endif
                    len = length;
                    return;
                }
                case raw_header_64:
                    type = qstype::RAW;
                    len = reader.template get_pod_contiguous<uint64_t>();
                    validate_object_length_64(len, "Raw vector length");
                    return;
                case attribute_header_8:
                    type = qstype::ATTRIBUTE;
                    len = reader.template get_pod_contiguous<uint8_t>();
                    return;
                case attribute_header_32:
                {
                    type = qstype::ATTRIBUTE;
                    const uint32_t length = reader.template get_pod_contiguous<uint32_t>();
                    validate_attr_length_32(length);
                    len = length;
                    return;
                }
                default:
                    reader.cleanup_and_throw("Unknown header type");
            }
        } else { // h5
            len = header_byte & bitmask_length_5;
            switch(header_byte_5) {
                case logical_header_5:
                    type = qstype::LOGICAL;
                    return;
                case integer_header_5:
                    type = qstype::INTEGER;
                    return;
                case numeric_header_5:
                    type = qstype::REAL;
                    return;
                case character_header_5:
                    type = qstype::CHARACTER;
                    return;
                case list_header_5:
                    type = qstype::LIST;
                    return;
                case attribute_header_5:
                    type = qstype::ATTRIBUTE;
                    return;
                default:
                    reader.cleanup_and_throw("Unknown header type");
            }
        }
    }
    public:

    // len, attr_length should be pre-initialized to 0
    void read_header(qstype & type, uint64_t & object_length, uint32_t & attr_length) {
        uint8_t header_byte = reader.template get_pod<uint8_t>();
        read_header_impl(header_byte, type, object_length);
        if(type == qstype::ATTRIBUTE) {
            attr_length = object_length;
            header_byte = reader.template get_pod_contiguous<uint8_t>();
            read_header_impl(header_byte, type, object_length);
            // if the first header_byte is attribute, the next header_byte must be a real object type
            if(type == qstype::ATTRIBUTE) {
                reader.cleanup_and_throw("Unknown header type");
            }
        }
    }

    inline void read_string_header(uint32_t & string_len) {
        string_len = reader.template get_pod<uint8_t>();
        switch(string_len) {
            case string_header_16: // 253
                string_len = reader.template get_pod_contiguous<uint16_t>();
                break;
            case string_header_32: // 254
            {
                const uint32_t length = reader.template get_pod_contiguous<uint32_t>();
                validate_string_length_32(length, "String length");
                string_len = length;
                break;
            }
            case string_header_NA: // 255
                string_len = NA_STRING_LENGTH;
                break;
            default:
                break;
        }
    }

    void read_and_assign_attributes(SEXP object, const uint32_t attr_length) {
        if(attr_length == 0) return;
#if R_VERSION >= R_Version(4, 6, 0)
        // assign attribute placeholder immediately for protection
        // symbol used doesn't matter as long it is not a specially intrepreted attribute symbol; it is just a placeholder
        // R_SpecSymbol is completely unused at all in r-source code, so it should be safe
        SEXP aptr = Rf_allocList(static_cast<int>(attr_length));
        Rf_setAttrib(object, R_SpecSymbol, aptr);
        std::string attr_name; // use std::string here, must be null terminated for Rf_install
        for(uint32_t i=0; i<attr_length; ++i) {
            uint32_t string_len;
            read_string_header(string_len);
            if(string_len == NA_STRING_LENGTH) {
                reader.cleanup_and_throw("Attribute names cannot be NA");
            }
            attr_name.resize(string_len);
            if(string_len > 0) {
                reader.get_data(attr_name.data(), string_len);
            }
            SET_TAG(aptr, Rf_install(attr_name.c_str()));
            SEXP aobj = read_object();
            SETCAR(aptr, aobj);
            aptr = CDR(aptr);
        }
        delayed_attributes.add(object);
#else
        SEXP aptr = Rf_allocList(static_cast<int>(attr_length));
        SET_ATTRIB(object, aptr); // assign immediately for protection
        bool set_class = false;
        SEXP class_attr = R_NilValue;
        std::string attr_name; // use std::string here, must be null terminated for Rf_install
        for(uint32_t i=0; i<attr_length; ++i) {
            uint32_t string_len;
            read_string_header(string_len);
            if(string_len == NA_STRING_LENGTH) {
                reader.cleanup_and_throw("Attribute names cannot be NA");
            }
            attr_name.resize(string_len);
            if(string_len > 0) {
                reader.get_data(attr_name.data(), string_len);
            }
            SET_TAG(aptr, Rf_install(attr_name.c_str()));
            SEXP aobj = read_object();
            SETCAR(aptr, aobj);
            aptr = CDR(aptr);
            
            // setting class name also sets object bit
            // SET_OBJECT is no longer part of API in R 4.5, so we use a workaround to set the object bit
            if( (strcmp(attr_name.c_str(), "class") == 0) && Rf_isString(aobj) && (Rf_xlength(aobj) >= 1) ) {
                set_class = true;
                class_attr = aobj;
            }
        }
        if(set_class) {
            Rf_setAttrib(object, R_ClassSymbol, class_attr); // sets the object bit, class names already assigned to attributes in the for loop
        }
#endif
    }

    SEXP read_object() {
        R_CheckStack();
        SEXP object = nullptr; // initialize to null value to hush compiler warnings
        qstype type;
        uint64_t object_length = 0;
        uint32_t attr_length = 0;
        read_header(type, object_length, attr_length);
        switch(type) {
            case qstype::NIL:
                return R_NilValue; // R_NilValue cannot have attributes, so return immediately
            case qstype::LOGICAL:
                object = PROTECT(Rf_allocVector(LGLSXP, static_cast<R_xlen_t>(object_length)));
                read_and_assign_attributes(object, attr_length);
                if(object_length > 0) integer_sexp.push_back(std::make_pair(object, object_length));
                // reader.get_data( reinterpret_cast<char*>(LOGICAL(object)), object_length*4);
                break;
            case qstype::INTEGER:
                object = PROTECT(Rf_allocVector(INTSXP, static_cast<R_xlen_t>(object_length)));
                read_and_assign_attributes(object, attr_length);
                if(object_length > 0) integer_sexp.push_back(std::make_pair(object, object_length));
                // reader.get_data( reinterpret_cast<char*>(INTEGER(object)), object_length*4 );
                break;
            case qstype::REAL:
                object = PROTECT(Rf_allocVector(REALSXP, static_cast<R_xlen_t>(object_length)));
                read_and_assign_attributes(object, attr_length);
                if(object_length > 0) real_sexp.push_back(std::make_pair(object, object_length));
                // reader.get_data( reinterpret_cast<char*>(REAL(object)), object_length*8 );
                break;
            case qstype::COMPLEX:
                object = PROTECT(Rf_allocVector(CPLXSXP, static_cast<R_xlen_t>(object_length)));
                read_and_assign_attributes(object, attr_length);
                if(object_length > 0) complex_sexp.push_back(std::make_pair(object, object_length));
                // reader.get_data( reinterpret_cast<char*>(COMPLEX(object)), object_length*16 );
                break;
            case qstype::CHARACTER:
            {
                object = PROTECT(Rf_allocVector(STRSXP, static_cast<R_xlen_t>(object_length)));
                read_and_assign_attributes(object, attr_length);
                if(object_length > 0) character_sexp.push_back(std::make_pair(object, object_length));
                break;
            }
            case qstype::LIST:
            {
                object = PROTECT(Rf_allocVector(VECSXP, static_cast<R_xlen_t>(object_length)));
                read_and_assign_attributes(object, attr_length);
                for(uint64_t i=0; i<object_length; ++i) {
                    SET_VECTOR_ELT(object, i, read_object());
                }
                break;
            }
            case qstype::RAW:
                object = PROTECT(Rf_allocVector(RAWSXP, static_cast<R_xlen_t>(object_length)));
                read_and_assign_attributes(object, attr_length);
                if(object_length > 0) raw_sexp.push_back(std::make_pair(object, object_length));
                // reader.get_data( reinterpret_cast<char*>(RAW(object)), object_length );
                break;
            default:
                // this statement should be unreachable
                reader.cleanup_and_throw("something went wrong (reading object type)");
        }
        UNPROTECT(1);
        return object;
    }

    SEXP read_root_object(uint64_t & runtime_hash) {
        SEXP object = PROTECT(read_object());
        try {
            read_object_data();
            reader.finish();
            runtime_hash = reader.get_hash_digest();
            UNPROTECT(1);
            return object;
        } catch (...) {
            UNPROTECT(1);
            throw;
        }
    }

    void read_object_data() {
        for(auto & x : character_sexp) {
            SEXP object = x.first;
            uint64_t object_length = x.second;
            for(uint64_t i=0; i<object_length; ++i) {
                uint32_t string_length;
                read_string_header(string_length);
                if(string_length == NA_STRING_LENGTH) {
                    SET_STRING_ELT(object, i, NA_STRING);
                } else if(string_length == 0) {
                    SET_STRING_ELT(object, i, R_BlankString);
                } else {
                    const char * string_ptr = reader.get_ptr(string_length);
                    if(string_ptr == nullptr) {
                        std::unique_ptr<char[]> string_buf(MAKE_UNIQUE_BLOCK(string_length));
                        reader.get_data(string_buf.get(), string_length);
                        SET_STRING_ELT(object, i, Rf_mkCharLenCE(string_buf.get(), static_cast<int>(string_length), CE_UTF8));
                    } else {
                        SET_STRING_ELT(object, i, Rf_mkCharLenCE(string_ptr, static_cast<int>(string_length), CE_UTF8));
                    }
                }
            }
        }
        for(auto & x : complex_sexp) {
            SEXP object = x.first;
            uint64_t object_length = x.second;
            reader.get_data( reinterpret_cast<char*>(COMPLEX(object)), object_length * 16 );
        }
        for(auto & x : real_sexp) {
            SEXP object = x.first;
            uint64_t object_length = x.second;
            reader.get_data( reinterpret_cast<char*>(REAL(object)), object_length * 8 );
        }
        for(auto & x : integer_sexp) {
            SEXP object = x.first;
            uint64_t object_length = x.second;
            reader.get_data( reinterpret_cast<char*>(INTEGER(object)), object_length * 4 );
        }
        for(auto & x : raw_sexp) {
            SEXP object = x.first;
            uint64_t object_length = x.second;
            reader.get_data( reinterpret_cast<char*>(RAW(object)), object_length );
        }

#if R_VERSION >= R_Version(4, 6, 0)
        delayed_attributes.resolve();
#endif
    }
};

#endif
