#include <fstream>
#include <corvusoft/restbed/session.hpp>
#include <corvusoft/restbed/request.hpp>
#include <corvusoft/restbed/response.hpp>
#include <spdlog/spdlog.h>
#include <stb_image.h>
#include <stb_image_resize2.h>
#include <stb_image_write.h>
#include "server_internal.hxx"

static constexpr std::size_t chunk_size = 1024*1024, max_image_size = 1024*1024*50, preview_size=480;
static const std::set<std::string> image_extension = {".png", ".jpg", ".jpeg", ".tga", ".bmp", ".psd", ".gif", ".jfif", ".pjpeg", ".pjp"};

struct UploadInfo {
    Server *server;
    std::shared_lock<std::shared_mutex> node_lock;
    std::size_t to_read;
    std::filesystem::path path;
    std::ofstream file;
    std::shared_ptr<Node> node;
};

void make_preview(const std::shared_ptr<UploadInfo>& info) {
    int x, y, channels;
    auto img = std::unique_ptr<stbi_uc, decltype(&free)>
        {stbi_load(info->path.c_str(), &x, &y, &channels, 0), &free};
    if (!img)
        return;

    float x_ration = (float)preview_size / (float)x, y_ration = (float)preview_size / (float)y;
    float ratio = std::min(x_ration, y_ration);
    int new_x = (int)((float)(x)*ratio), new_y = (int)((float)(y)*ratio);
    stbir_pixel_layout layout;
    switch (channels) {
        case 1: layout = STBIR_1CHANNEL; break;
        case 2: layout = STBIR_2CHANNEL; break;
        case 3: layout = STBIR_RGB; break;
        case 4: layout = STBIR_RGBA; break;
        default: return;
    }

    auto rimg = std::unique_ptr<unsigned char, decltype(&free)>
        {stbir_resize_uint8_linear(img.get(), x, y, 0, nullptr, new_x, new_y, 0, layout), &free};
    if (!rimg)
        return;

    auto png_path = info->path.replace_extension("png");
    if (!stbi_write_png(png_path.c_str(), new_x, new_y, channels, rimg.get(), 0))
        return;

    info->node->preview = true;
}

void fetch_handler(const std::shared_ptr<restbed::Session> &s, const restbed::Bytes &bytes) {
    std::shared_ptr<UploadInfo> info = s->get("upload");

    std::size_t read = bytes.size();
    info->to_read -= std::min(read, info->to_read);
    info->file.write((char*)bytes.data(), bytes.size());
    if (info->to_read > 0)
        return s->fetch(std::min(info->to_read, chunk_size), fetch_handler);
    info->file.close();
    s->close(200);

    std::size_t real_size = std::filesystem::file_size(info->path);
    info->node->size = real_size;
    auto ext = std::filesystem::path{info->node->name}.extension().string();
    if (real_size < max_image_size && image_extension.contains(ext))
        make_preview(info);

    info->node_lock.unlock();
    info->server->save();
}

void Server::upload(const std::shared_ptr<restbed::Session> &s) {
    const auto req = s->get_request();
    if (!req->has_header("X-Node"))
        return s->close(400, "Missing node");
    if (!req->has_header("X-Token"))
        return s->close(400, "Missing token");

    if (req->get_header("Transfer-Encoding") == "chunked") {
        spdlog::error("Encountered a chunked upload!");
        return s->close(500, "Sorry but your browser is not supported yet");
    }

    std::uint64_t node_id = req->get_header("X-Node", 0);
    std::string token = req->get_header("X-Token");

    check_user() return s->close(400, "Invalid user");
    {
        std::shared_lock lock{user->node_lock};
        auto node = get_node(user, node_id);
        if (!node) return s->close(400, "Invalid node");
        if (!node->file) return s->close(400, "Can't upload to a directory");
        std::size_t to_read = req->get_header("Content-Length", 0);
        auto path = user->user_dir / std::to_string(node->id);
        if (node->preview) {
            node->preview = false;
            std::filesystem::remove(path.replace_extension("png"));
        }
        std::shared_ptr<UploadInfo> info{new UploadInfo{
            .server = this,
            .node_lock = std::shared_lock{user->node_lock},
            .to_read = to_read,
            .path = path,
            .file = std::ofstream{path, std::ios_base::out|std::ios_base::trunc|std::ios_base::binary},
            .node = node
        }};
        s->set("upload", info);
        s->fetch(std::min(to_read, chunk_size), fetch_handler);
    }
}