/******************************************************************************
 Copyright 2020 Hypersect LLC
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
 
     http://www.apache.org/licenses/LICENSE-2.0
 
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
******************************************************************************/

/******************************************************************************
 RJ_IMAGE SINGLE FILE HEADER LIBRARY USAGE
   This file contains both the header and implementation. By default,
   including this file will function like a header. To compile the
   implementation, include this file from one source file with one of the 
   following preprocessor definitions prior to it.
   
   For little endian systems:
    #define RJIMAGE_IMPLEMENTATION
    #define RJIMAGE_ENDIAN RJIMAGE_ENDIAN_LITTLE
    #include "RJImage.h"
   
   For big endian systems:
    #define RJIMAGE_IMPLEMENTATION
    #define RJIMAGE_ENDIAN RJIMAGE_ENDIAN_BIG
    #include "RJImage.h"

 RJ_IMAGE FILE FORMAT SPECIFICATION
   All data is little endian
   
   The file starts with a header data block
    c8  tag[8]                           // 'RJ_IMAGE'
    u32 version                          // 1
    u32 dimension_count                  // number of dimensions (1D, 2D, etc)
    u32 channel_count                    // number of channels per pixel
    u32 channel_type                     // numerical format of channel bits
    u32 channel_bytes                    // bytes of data per channel
    u32 dimension_sizes[dimension_count] // pixels per dimension
    u32 channel_usages[channel_count]    // what channels represnt
   
   The header block is immediately followed by an image data block.
   Pixels are packed as nested arrays with lower dimensions contiguous. For
   example, a 3D texture would be packed as an array of pixels[z][x][y].
   Each pixel is represented as array of interleaved channel data.
******************************************************************************/

#ifndef RJIMAGE_INCLUDED_HEADER
#define RJIMAGE_INCLUDED_HEADER

// Include header for standard types
#include <cstdint>

// Channel types describe how to interpret the pixel channel bits as a number.
// Values below 0x10000000 are reserved for the file format.
#define RJIMAGE_CHANNEL_TYPE_UNSIGNED 0x00000000
#define RJIMAGE_CHANNEL_TYPE_SIGNED   0x00000001
#define RJIMAGE_CHANNEL_TYPE_FLOAT    0x00000002
#define RJIMAGE_CHANNEL_TYPE_USER     0x10000000

// Channel usages describe what pixel channel values represent.
// Values below 0x10000000 are reserved for the file format.
#define RJIMAGE_CHANNEL_USAGE_NONE  0x00000000
#define RJIMAGE_CHANNEL_USAGE_RED   0x00000001
#define RJIMAGE_CHANNEL_USAGE_GREEN 0x00000002
#define RJIMAGE_CHANNEL_USAGE_BLUE  0x00000003
#define RJIMAGE_CHANNEL_USAGE_ALPHA 0x00000004
#define RJIMAGE_CHANNEL_USAGE_USER  0x10000000 // User defined range starts here

// Description of data contained in an image file.
struct RJImage_Descriptor
{
    uint32_t dimension_count;
    uint32_t channel_count;
    uint32_t channel_type;
    uint32_t channel_bytes;

    uint32_t dimension_sizes[16]; // dimension_count array elements are valid
    uint32_t channel_usages[16];  // channel_count array elements are valid
};

// Read file data from a memory buffer into an image descriptor. The file data
// buffer is not assumed to be from a safe source and is validated accordingly.
// On success (the file format is valid), a pointer to the little-endian pixel
// data block is returned. 
// On failure (the file format is invalid), nullptr is returned.
const void* RJImage_ReadFileHeaderFromMemory(RJImage_Descriptor* out_descriptor, const void* file_data, size_t file_size);

// Write file with a given descriptor and pixel array into a file buffer.
// It is assumed that the file data is large enough to store the necessary 
// file header data. It is assumed that the descriptor data is from a safe 
// source and math does not need to be tested for overflow. A pointer to where
// the little-endian pixel data block should be written is returned.
void* RJImage_WriteFileHeaderToMemory(void* out_file_data, const RJImage_Descriptor& descriptor);

// Get the byte size of the header data block for a given image descriptor.
// It is assumed that the descriptor is from a safe source and math overflow
// will not occur.
size_t RJImage_GetHeaderSize(const RJImage_Descriptor& descriptor);

// Get the byte size of pixel data block for a given image descriptor.
// It is assumed that the descriptor is from a safe source and math overflow
// will not occur.
size_t RJImage_GetPixelDataSize(const RJImage_Descriptor& descriptor);

// Get the byte size needed to store an entire file given its image descriptor.
// It is assumed that the descriptor is from a safe source and math overflow will not occur.
size_t RJImage_GetFileSize(const RJImage_Descriptor& descriptor);

#endif // RJIMAGE_INCLUDED_HEADER

//*****************************************************************************
//*****************************************************************************
// END OF HEADER CODE AND START OF IMPLEMENTATION CODE
//*****************************************************************************
//*****************************************************************************

#ifdef RJIMAGE_IMPLEMENTATION

#define RJIMAGE_ENDIAN_LITTLE 1
#define RJIMAGE_ENDIAN_BIG    2

#if RJIMAGE_ENDIAN == RJIMAGE_ENDIAN_LITTLE
inline uint32_t RJImage_ReadU32(const void* data)
{
    return *(const uint32_t*)data;
}

inline void RJImage_WriteU32(void* data, uint32_t value)
{
    *(uint32_t*)data = value;
}
#elif RJIMAGE_ENDIAN == RJIMAGE_ENDIAN_BIG
inline uint32_t RJImage_EndianSwapU32(uint32_t value) 
{
    return ((value >> 24) & 0x000000FF) |
           ((value >>  8) & 0x0000FF00) |
           ((value <<  8) & 0x00FF0000) |
           ((value << 24) & 0xFF000000);
}

inline uint32_t RJImage_ReadU32(const void* data)
{
    return RJImage_EndianSwapU32( *(const uint32_t*)data );
}

inline void RJImage_WriteU32(void* data, uint32_t value) {
    *(uint32_t*)data = RJImage_EndianSwapU32(value);
}
#endif

const void* RJImage_ReadFileHeaderFromMemory(RJImage_Descriptor* out_descriptor, const void* file_data, size_t file_size)
{
    const uint8_t* file_bytes = (const uint8_t*)file_data;

    // only support files up to 4gb to similify testing for overflow
    const uint64_t max_file_size = 0xFFFFFFFF;
    if (file_size > max_file_size)
        return nullptr;

    // read header
    const size_t static_header_size = 28;
    if (file_size < static_header_size)
        return nullptr;

    const uint8_t* tag = file_bytes + 0;
    if (tag[0] != 'R' || 
        tag[1] != 'J' || 
        tag[2] != '_' || 
        tag[3] != 'I' || 
        tag[4] != 'M' || 
        tag[5] != 'A' || 
        tag[6] != 'G' || 
        tag[7] != 'E')
        return nullptr;

    uint32_t version         = RJImage_ReadU32(file_bytes+8);
    uint32_t dimension_count = RJImage_ReadU32(file_bytes+12);
    uint32_t channel_count   = RJImage_ReadU32(file_bytes+16);
    uint32_t channel_type    = RJImage_ReadU32(file_bytes+20);
    uint32_t channel_bytes   = RJImage_ReadU32(file_bytes+24);

    if (version != 1)
        return nullptr;

    // Verify we fit within limitations of descriptor. These aren't actual
	// limitations  of the file format, but it keeps this helper API simple.
    if (dimension_count > sizeof(out_descriptor->dimension_sizes)/sizeof(out_descriptor->dimension_sizes[0]))
        return nullptr;
    
    if (channel_count > sizeof(out_descriptor->channel_usages)/sizeof(out_descriptor->channel_usages[0]))
        return nullptr;

    // Check that there is room for the non-static header data
    uint64_t dimension_data_size = (uint64_t)dimension_count*4;
    uint64_t channel_data_size   = (uint64_t)channel_count*4;
    uint64_t total_header_size   = static_header_size + dimension_data_size + channel_data_size;
    if (total_header_size > file_size)
        return nullptr;

    // Read the non-static header data
    const uint32_t* file_dimension_sizes = (uint32_t*)(file_bytes+static_header_size);
    const uint32_t* file_channel_usages  = (uint32_t*)(file_bytes+static_header_size+(size_t)dimension_data_size);

    for (size_t i = 0; i < dimension_count; ++i)
        out_descriptor->dimension_sizes[i] = RJImage_ReadU32(file_dimension_sizes+i);

    for (size_t i = 0; i < channel_count; ++i)
        out_descriptor->channel_usages[i] = RJImage_ReadU32(file_channel_usages+i);

    // Compute the pixel data size
    uint64_t pixel_data_size = 0;
    if (dimension_count > 0)
    {
        pixel_data_size = channel_count * channel_bytes;
        if (pixel_data_size > file_size) // check data size before overflow can occur
             return nullptr;

        for (size_t i = 0; i < dimension_count; ++i)
        {
            pixel_data_size *= out_descriptor->dimension_sizes[i];
            if (pixel_data_size > file_size) // check data size before overflow can occur
                return nullptr;
        }
    }

    // Check that there is enough space for the pixel data
    if (total_header_size + pixel_data_size > file_size)
        return nullptr;

    const void* pixel_data = file_bytes + total_header_size;

    // fill in remaining descriptor data and return success
    out_descriptor->dimension_count = dimension_count;
    out_descriptor->channel_count   = channel_count;
    out_descriptor->channel_type    = channel_type;
    out_descriptor->channel_bytes   = channel_bytes;

    return pixel_data;
}

void* RJImage_WriteFileHeaderToMemory(void* out_file_data, const RJImage_Descriptor& descriptor)
{
    uint8_t* file_bytes = (uint8_t*)out_file_data;

    const size_t static_header_size = 28;
    const uint32_t version = 1;

    uint32_t dimension_count = descriptor.dimension_count;
    uint32_t channel_count   = descriptor.channel_count;

    // Write the static header data
    uint8_t* tag = file_bytes + 0;
    tag[0] = 'R'; 
    tag[1] = 'J'; 
    tag[2] = '_'; 
    tag[3] = 'I'; 
    tag[4] = 'M'; 
    tag[5] = 'A'; 
    tag[6] = 'G'; 
    tag[7] = 'E';

    RJImage_WriteU32(file_bytes+8,  version);
    RJImage_WriteU32(file_bytes+12, dimension_count);
    RJImage_WriteU32(file_bytes+16, channel_count  );
    RJImage_WriteU32(file_bytes+20, descriptor.channel_type   );
    RJImage_WriteU32(file_bytes+24, descriptor.channel_bytes  );

    // Write the non-static header data
    size_t dimension_data_size = (size_t)dimension_count*4;
    size_t channel_data_size   = (size_t)channel_count*4;
    uint32_t* file_dimension_sizes = (uint32_t*)(file_bytes+static_header_size);
    uint32_t* file_channel_usages  = (uint32_t*)(file_bytes+static_header_size+dimension_data_size);

    for (size_t i = 0; i < dimension_count; ++i)
         RJImage_WriteU32(file_dimension_sizes+i, descriptor.dimension_sizes[i]);

    for (size_t i = 0; i < channel_count; ++i)
         RJImage_WriteU32(file_channel_usages+i, descriptor.channel_usages[i]);

    // Return the pixel data location
    void* pixel_data = file_bytes + static_header_size + dimension_data_size + channel_data_size;
    return pixel_data;
}

size_t RJImage_GetHeaderSize(const RJImage_Descriptor& descriptor)
{
    const size_t static_header_size = 28;
    size_t dimension_data_size = (size_t)descriptor.dimension_count*4;
    size_t channel_data_size   = (size_t)descriptor.channel_count*4;
    
    size_t total_header_size = static_header_size + dimension_data_size + channel_data_size;
    return total_header_size;
}

size_t RJImage_GetPixelDataSize(const RJImage_Descriptor& descriptor)
{
    size_t pixel_data_size = 0;
    if (descriptor.dimension_count > 0)
    {
        pixel_data_size = descriptor.channel_count * descriptor.channel_bytes;
        for (size_t i = 0; i < descriptor.dimension_count; ++i)
            pixel_data_size *= descriptor.dimension_sizes[i];
    }
    return pixel_data_size;
}

size_t RJImage_GetFileSize(const RJImage_Descriptor& descriptor)
{
    size_t header_size     = RJImage_GetHeaderSize(descriptor);
    size_t pixel_data_size = RJImage_GetPixelDataSize(descriptor);
    
    return header_size + pixel_data_size;
}

#endif