#include <cstdint>
#include <exception>
#include <filesystem>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>

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

extern "C" void to_grayscale(const uint32_t *in, uint8_t *out, uint32_t width, uint32_t height);
extern "C" void detect_edges(const uint8_t *in, uint8_t *out, uint32_t width, uint32_t height);

using RGBAImage = std::vector<uint32_t>;
using GrayscaleImage = std::vector<uint8_t>;

namespace {

std::filesystem::path buildOutputPath(const std::filesystem::path &inputPath, const std::string &op) {
    auto stem = inputPath.stem().string();
    if (stem.empty()) {
        stem = "output";
    }

    auto parent = inputPath.parent_path();
    auto outputName = stem + "_" + op + ".png";
    return parent / outputName;
}

RGBAImage toRGBA(stbi_uc* pixels, int channels, uint32_t width, uint32_t height) {
    const auto numPixels = width * height;
    RGBAImage rgbaPixels(width * height, 0);
    switch(channels) {
        case 3: {
            for(size_t idx = 0; idx < numPixels; ++idx) {
                const auto r = pixels[idx * 3];
                const auto g = pixels[idx * 3 + 1];
                const auto b = pixels[idx * 3 + 2];
                rgbaPixels[idx] = r | (g << 8) | (b << 16) | 0xff000000;
            }
            break;
        }
        case 4: {
            std::memcpy(reinterpret_cast<uint8_t*>(rgbaPixels.data()), pixels, numPixels * 4); 
            break;
        }
        default:
            throw std::runtime_error{"Invalid number of channels in image, only RGB and RGBA are supported!"};
    }

    return rgbaPixels;
}
}

int main(int argc, char **argv) try {
    if (argc != 2) {
        std::cerr << "Usage: " << argv[0] << " <image-path>" << '\n';
        return EXIT_FAILURE;
    }

    const std::filesystem::path inputPath = argv[1];
    if (!std::filesystem::exists(inputPath)) {
        std::cerr << "Input file does not exist: " << inputPath << '\n';
        return EXIT_FAILURE;
    }

    int width = 0;
    int height = 0;
    int originalChannels = 0;
    stbi_uc *pixels = stbi_load(inputPath.string().c_str(), &width, &height, &originalChannels, 0);
    if (!pixels) {
        std::cerr << "Failed to load image: " << stbi_failure_reason() << '\n';
        return EXIT_FAILURE;
    }

    const uint32_t uWidth = static_cast<uint32_t>(width);
    const uint32_t uHeight = static_cast<uint32_t>(height);
    const size_t pixelCount = static_cast<size_t>(uWidth) * static_cast<size_t>(uHeight);

    // Convert to RGBA so that everything is 32-bit aligned!
    const auto rgbaPixels = toRGBA(pixels, originalChannels, uWidth, uHeight);

    GrayscaleImage grayscalePixels(pixelCount, 0);
    to_grayscale(rgbaPixels.data(), grayscalePixels.data(), uWidth, uHeight);

    {
        const auto outputPath = buildOutputPath(inputPath, "grayscale");
        if (!stbi_write_png(outputPath.string().c_str(), uWidth, uHeight, 1, grayscalePixels.data(), static_cast<int>(uWidth))) {
            std::cerr << "Failed to write image: " << outputPath << '\n';
            stbi_image_free(pixels);
            return EXIT_FAILURE;
        }
        std::cout << "Saved grayscale output to " << outputPath << '\n';
    }

    GrayscaleImage edgeDetectionPixels(pixelCount, 0);
    detect_edges(grayscalePixels.data(), edgeDetectionPixels.data(), uWidth, uHeight);

    {
        const auto outputPath = buildOutputPath(inputPath, "edges");
        if (!stbi_write_png(outputPath.string().c_str(), uWidth, uHeight, 1, edgeDetectionPixels.data(), static_cast<int>(uWidth))) {
            std::cerr << "Failed to write image: " << outputPath << '\n';
            stbi_image_free(pixels);
            return EXIT_FAILURE;
        }
        std::cout << "Saved edge detection output to " << outputPath << '\n';
    }

    stbi_image_free(pixels);
    
    return EXIT_SUCCESS;
} catch (const std::exception &ex) {
    std::cerr << "Error: " << ex.what() << '\n';
    return EXIT_FAILURE;
}