#include <fstream>
#include <algorithm>
#include <ReImage/Image.h>

#define STB_IMAGE_IMPLEMENTATION
#include <stb_image.h>
#define STB_IMAGE_RESIZE_IMPLEMENTATION
#include <stb_image_resize2.h>

namespace ReImage {

    Image::Image(const std::string &file, const TextureFlag& flags) {
        load(file);
        if (flags& TextureFlag::INVERT_Y)
            invertY();

        this->flags = flags;
    }

    Image::Image(const Color4* data, const size_t& length, unsigned int width, unsigned int height, TextureFlag flags) {
        pixel_data.resize(length * 4);
        memcpy(pixel_data.data(), data, length * sizeof(Color4));
        this->width = width;
        this->height = height;
        this->format = TextureFormat::RGBA;
        if (flags & INVERT_Y)
            invertY();

        this->flags = flags;
    }


    void Image::load(const std::string& file) {
        std::ifstream ifStream(file, std::ios::in | std::ios::binary);
        unsigned char bmpFileHeader[14];

        ifStream.read(reinterpret_cast<char *>(bmpFileHeader), 14);
        if (bmpFileHeader[0] == 'B' && bmpFileHeader[1] == 'M') {
            ifStream.close();
            loadBMP(file);
            return;
        }

        //TODO determine if it's a PNG using the file header instead.
        if (file.ends_with(".png"))
            loadPNG(file);
    }

    void Image::loadBMP(const std::string &file) {
        std::ifstream bmpFile(file, std::ios::in | std::ios::binary);
        if (!bmpFile.is_open())
            return;

        unsigned char bmpFileHeader[14];
        unsigned char bmpInfoHeader[40];

        bmpFile.read(reinterpret_cast<char *>(bmpFileHeader), 14);
        bmpFile.read(reinterpret_cast<char *>(bmpInfoHeader), 40);

        width = bmpInfoHeader[4] + (bmpInfoHeader[5] << 8) + (bmpInfoHeader[6] << 16) + (bmpInfoHeader[7] << 24);
        height = bmpInfoHeader[8] + (bmpInfoHeader[9] << 8) + (bmpInfoHeader[10] << 16) + (bmpInfoHeader[11] << 24);

        int rowPadded = (width * 3 + 3) & (~3);
        pixel_data.resize(width * height * 3);

        std::vector<unsigned char> rowData(rowPadded);
        for (int y = height - 1; y >= 0; --y) {
            bmpFile.read(reinterpret_cast<char *>(rowData.data()), rowPadded);
            for (int x = 0; x < width; ++x) {
                pixel_data[(y * width + x) * 3 + 2] = rowData[x * 3 + 0];
                pixel_data[(y * width + x) * 3 + 1] = rowData[x * 3 + 1];
                pixel_data[(y * width + x) * 3 + 0] = rowData[x * 3 + 2];
            }
        }
        bmpFile.close();
        format = TextureFormat::RGB;
    }

    void Image::loadPNG(const std::string& file) {
        int channels, w, h;

        unsigned char* imageData = stbi_load(file.c_str(), &w, &h, &channels, 0);

        if (imageData == nullptr)
            return;

        width = w;
        height = h;
        if (channels == 3)
            format = TextureFormat::RGB;

        if (channels == 4)
            format = TextureFormat::RGBA;

        pixel_data.assign(imageData, imageData + width * height * channels);
        stbi_image_free(imageData);
    }

    void Image::invertY() {
        unsigned int rowSize;
        if (format == TextureFormat::RGB)
            rowSize = width * 3;
        if (format == TextureFormat::RGBA)
            rowSize = width * 4;

        std::vector<unsigned char> temp(rowSize);
        for (unsigned int y = 0; y < height / 2; ++y) {
            unsigned char *topRow = &pixel_data[y * rowSize];
            unsigned char *bottomRow = &pixel_data[(height - y - 1) * rowSize];
            std::copy(bottomRow, bottomRow + rowSize, temp.begin());
            std::copy(topRow, topRow + rowSize, bottomRow);
            std::copy(temp.begin(), temp.end(), topRow);
        }
    }

    unsigned int Image::getWidth() const {
        return width;
    }

    unsigned int Image::getHeight() const {
        return height;
    }

    TextureFormat Image::getTextureFormat() {
        return format;
    }

    TextureFlag Image::getFlags() {
        return flags;
    }

    Image Image::downscale(unsigned int rhs) {
        std::vector<unsigned char> result;

        if (format == TextureFormat::RGB)
            result.resize((width / rhs) * (height / rhs) * 3),
            stbir_resize_uint8_linear(pixel_data.data(), width, height, 0, result.data(), (width / rhs), (height / rhs), 0, stbir_pixel_layout::STBIR_RGB);

        else if (format == TextureFormat::RGBA)
            result.resize((width / rhs) * (height / rhs) * 4),
            stbir_resize_uint8_linear(pixel_data.data(), width, height, 0, result.data(), (width / rhs), (height / rhs), 0, stbir_pixel_layout::STBIR_RGBA);

        return Image(result, format, (width / rhs), (height / rhs));
    }

    Image::Image(std::vector<unsigned char>& pixel_data, TextureFormat format, unsigned int width, unsigned int height) {
        this->pixel_data = pixel_data;
        this->width = width;
        this->height = height;
        this->format = format;
        this->flags = TextureFlag::NONE;

    }

    Image::Image(const Color3* data, const size_t& length, unsigned int width, unsigned int height, TextureFlag flags) {
        pixel_data.resize(length * 3);
        memcpy(pixel_data.data(), data, length * sizeof(Color3));
        this->width = width;
        this->height = height;
        this->format = TextureFormat::RGB;
        if (flags & INVERT_Y)
            invertY();

        this->flags = flags;
    }
}