#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

#include "pipeline.h"
#include <hip/hip_runtime.h>
#include <cstdio>
#include <string>
#include <vector>

static float4* upload_image(const char* path, int& out_width, int& out_height) {
    int w, h, channels;
    unsigned char* data = stbi_load(path, &w, &h, &channels, 4);
    if (!data) throw std::runtime_error(std::string("Failed to load: ") + path);

    out_width = w; out_height = h;
    const int n = w * h;

    std::vector<float4> host(n);
    for (int i = 0; i < n; ++i) {
        host[i] = make_float4(
            data[4*i + 0] / 255.0f,
            data[4*i + 1] / 255.0f,
            data[4*i + 2] / 255.0f,
            0.0f);
    }
    stbi_image_free(data);

    float4* dev;
    hip_check(hipMalloc(&dev, sizeof(float4) * n));
    hip_check(hipMemcpy(dev, host.data(), sizeof(float4) * n, hipMemcpyHostToDevice));
    return dev;
}

static void dump_level(const PyramidLevel& lvl, const char* path) {
    const int n = lvl.width * lvl.height;
    std::vector<float4> host(n);
    hip_check(hipMemcpy(host.data(), lvl.color_smax,
                        sizeof(float4) * n, hipMemcpyDeviceToHost));

    std::vector<unsigned char> bytes(n * 3);
    for (int i = 0; i < n; ++i) {
        bytes[3*i + 0] = static_cast<unsigned char>(
            std::min(host[i].x, 1.0f) * 255.0f);
        bytes[3*i + 1] = static_cast<unsigned char>(
            std::min(host[i].y, 1.0f) * 255.0f);
        bytes[3*i + 2] = static_cast<unsigned char>(
            std::min(host[i].z, 1.0f) * 255.0f);
    }
    stbi_write_png(path, lvl.width, lvl.height, 3, bytes.data(), lvl.width * 3);
    printf("Wrote %s (%dx%d)\n", path, lvl.width, lvl.height);
}

int main(int argc, char* argv[]) {
    if (argc < 2) {
        fprintf(stderr, "Usage: %s <image>\n", argv[0]);
        return 1;
    }

    try {
        int w, h;
        float4* src = upload_image(argv[1], w, h);

        const int num_levels = 5;
        Pyramid pyr = make_pyramid(w, h, num_levels);

        hip_check(hipMemcpy(pyr.levels[0].color_smax, src,
                            sizeof(float4) * w * h, hipMemcpyDeviceToDevice));
        hipFree(src);

        build_pyramid(pyr);

        for (int i = 0; i < static_cast<int>(pyr.levels.size()); ++i) {
            std::string path = "level_" + std::to_string(i) + ".png";
            dump_level(pyr.levels[i], path.c_str());
        }

        destroy_pyramid(pyr);

    } catch (const std::exception& e) {
        fprintf(stderr, "Error: %s\n", e.what());
        return 1;
    }

    return 0;
}