#include "pch.h"
#include "log.h"
#include "image.h"
#include "asset.h"

#include <stb/stb_image.h>
#include <stb/stb_image_write.h>

bool Image::load(std::string filename)
{
    if (Asset::is_asset(filename))
    {
        Asset file;
        if (!(file.open(filename.c_str()) && file.read_all()))
        {
            file.close();
            return false;
        }

        std::string name, base, ext;
        std::regex r(R"((.*)[\\/]([^\\/]+)\.(\w+)$)");
        std::smatch m;
        if (!std::regex_search(filename, m, r))
            return false;
        file_base = m[1].str();
        file_name = m[2].str();
        file_ext = m[3].str();
        
        stbi_set_flip_vertically_on_load(false);
        uint8_t* buffer = stbi_load_from_memory(file.m_data, file.m_len, &width, &height, nullptr, 4);
        file.close();
        comp = 4;
        m_data = std::unique_ptr<uint8_t[]>(buffer);
        return true;
    }
    else
    {
        return load_file(filename);
    }
}

bool Image::load_file(std::string filename)
{
    stbi_set_flip_vertically_on_load(false);
    uint8_t* buffer = stbi_load(filename.c_str(), &width, &height, nullptr, 4);
    if (!buffer)
        return false;
    comp = 4;
    m_data = std::unique_ptr<uint8_t[]>(buffer);

    std::string name, base, ext;
    std::regex r(R"((.*)[\\/]([^\\/]+)\.(\w+)$)");
    std::smatch m;
    if (!std::regex_search(filename, m, r))
        return false;
    file_base = m[1].str();
    file_name = m[2].str();
    file_ext = m[3].str();

    return true;
}

bool Image::save_png(const std::string& path) const noexcept
{
    bool ret = stbi_write_png(path.c_str(), width, height, comp, data(), 0);
    if (!ret)
        LOG("failed Image::save_png %s", path.c_str());
    return ret;
}

bool Image::save_jpg(const std::string& path, int quality) const noexcept
{
    bool ret = stbi_write_jpg(path.c_str(), width, height, comp, data(), quality);
    if (!ret)
        LOG("failed Image::save_jpg %s", path.c_str());
    return ret;
}

void Image::flip()
{
    auto flipped = std::make_unique<uint8_t[]>(width*height*4);
    int line_size = width * 4;
    const uint8_t* src = m_data.get();
    uint8_t* dst = flipped.get() + line_size * (height - 1);
    for (int y = 0; y < height; y++)
    {
        std::copy(src, src+line_size, dst);
        src += line_size;
        dst -= line_size;
    }
    std::swap(m_data, flipped);
}

void Image::gayscale_alpha()
{
    int np = width * height;
    auto ptr = reinterpret_cast<glm::u8vec4*>(m_data.get());
    for (int i = 0; i < np; i++)
    {
        auto& c = ptr[i];
        c.g = c.b = c.r = 255 - c.a;
        c.a = 255;
    }
}

Image Image::resize(int w, int h) const
{
    Image ret;
    ret.create(w, h);
    ret.file_base = file_base;
    ret.file_name = file_name;
    ret.file_ext = file_ext;
    auto temp = (glm::u8vec4*)ret.data();
    auto pixels = (glm::u8vec4*)data();
    float x_ratio = ((float)(width - 1)) / w;
    float y_ratio = ((float)(height - 1)) / h;
    int offset = 0;
    for (int i = 0; i < h; i++) {
        for (int j = 0; j < w; j++) {
            int x = (int)(x_ratio * j);
            int y = (int)(y_ratio * i);
            float x_diff = (x_ratio * j) - x;
            float y_diff = (y_ratio * i) - y;
            int index = y * width + x;

            // range is 0 to 255 thus bitwise AND with 0xff
            glm::vec4 A = pixels[index];
            glm::vec4 B = pixels[index + 1];
            glm::vec4 C = pixels[index + width];
            glm::vec4 D = pixels[index + width + 1];

            // Y = A(1-w)(1-h) + B(w)(1-h) + C(h)(1-w) + Dwh
            glm::vec4 gray = A*(1 - x_diff)*(1 - y_diff) + B * (x_diff)*(1 - y_diff) +
                C * (y_diff)*(1 - x_diff) + D * (x_diff*y_diff);
            temp[offset++] = glm::clamp(gray, glm::vec4(0), glm::vec4(255));
        }
    }
    return ret;
}

Image Image::resize_power2() const
{
    int w = pow(2, ceil(log(width) / log(2)));
    int h = pow(2, ceil(log(height) / log(2)));
    if (w == width && h == height)
    {
        Image i;
        i.create(width, height);
        i.file_base = file_base;
        i.file_name = file_name;
        i.file_ext = file_ext;
        std::copy(m_data.get(), m_data.get() + width * height * comp, i.m_data.get());
        return i;
    }
    return resize(w, h);
}

Image Image::resize_squared(const glm::u8vec4& bg) const
{
    Image ret;
    if (width == height)
    {
        ret.create(width, height);
        ret.file_base = file_base;
        ret.file_name = file_name;
        ret.file_ext = file_ext;
        std::copy(m_data.get(), m_data.get() + width * height * comp, ret.m_data.get());
    }
    else
    {
        int pad_x = 0;
        int pad_y = 0;
        int size = 0;
        if (height > width)
        {
            size = height;
            pad_x = (size - width) / 2;
        }
        else
        {
            size = width;
            pad_y = (size - height) / 2;
        }
        ret.create(size, size);
        ret.file_base = file_base;
        ret.file_name = file_name;
        ret.file_ext = file_ext;
        auto ptr_src = reinterpret_cast<glm::u8vec4*>(m_data.get());
        auto ptr_dst = reinterpret_cast<glm::u8vec4*>(ret.m_data.get());
        std::fill_n(ptr_dst, size * size, bg);
        for (int y = 0; y < height; y++)
            std::copy_n(ptr_src + y * width, width, ptr_dst + pad_x + (y + pad_y) * ret.width);
    }
    return ret;
}

bool Image::read(BinaryStreamReader& r)
{
    Serializer::Descriptor d;
    r >> d;
    if (d.class_id != "image_png")
        return false;
    d.value<Serializer::Integer>("width", width);
    d.value<Serializer::Integer>("height", height);
    d.value<Serializer::Integer>("comp", comp);
    file_base = wstr2str(d.value<Serializer::String>("file_base"));
    file_name = wstr2str(d.value<Serializer::String>("file_name"));
    file_ext = wstr2str(d.value<Serializer::String>("file_ext"));
    auto img_raw = d.get<Serializer::RawData>("data");
    int png_width, png_height, png_comp;
    m_data = std::unique_ptr<uint8_t[]>(stbi_load_from_memory(
        img_raw->data.data(), img_raw->data.size(), &png_width, &png_height, &png_comp, 4));
    return true;
}

void Image::write(BinaryStreamWriter& w) const
{
    Serializer::Descriptor d;
    d.class_id = "image_png";
    d.name = L"Image class";

    d.props["width"] = std::make_shared<Serializer::Integer>(width);
    d.props["height"] = std::make_shared<Serializer::Integer>(height);
    d.props["comp"] = std::make_shared<Serializer::Integer>(comp);
    d.props["file_base"] = std::make_shared<Serializer::String>(file_base);
    d.props["file_name"] = std::make_shared<Serializer::String>(file_name);
    d.props["file_ext"] = std::make_shared<Serializer::String>(file_ext);

    // really ugly way to compress the png and store it with a lambda
    stbi_write_png_to_func([](void* context, void* data, int size) {
        Serializer::Descriptor& d = *static_cast<Serializer::Descriptor*>(context);
        d.props["data"] = std::make_shared<Serializer::RawData>(std::vector<uint8_t>((uint8_t*)data, (uint8_t*)data + size));
    }, &d, width, height, comp, m_data.get(), 0);

    w << d;
}