image.cpp

#include <math.h>
#ifndef INFINITY
#define INFINITY __builtin_huge_valf()
#endif
#include <Foundation/Foundation.hpp>
#include <Metal/Metal.hpp>

#include "backend/metal/mt_api.h"
#include "core/logger.h"

namespace rhi {

void destroy_image_views(Device &device, Image &image);

static MTL::PixelFormat to_pixel_format(rhi::ImageFormat format) {
    switch (format) {
    case rhi::ImageFormat::R8_UNORM:
        return MTL::PixelFormatR8Unorm;
    case rhi::ImageFormat::RGBA8_UNORM:
        return MTL::PixelFormatRGBA8Unorm;
    case rhi::ImageFormat::RGBA8_SRGB:
        return MTL::PixelFormatRGBA8Unorm_sRGB;
    case rhi::ImageFormat::BGRA8_UNORM:
        return MTL::PixelFormatBGRA8Unorm;
    case rhi::ImageFormat::RGBA16_FLOAT:
        return MTL::PixelFormatRGBA16Float;
    case rhi::ImageFormat::R32_FLOAT:
        return MTL::PixelFormatR32Float;
    case rhi::ImageFormat::D32_FLOAT:
        return MTL::PixelFormatDepth32Float;
    default:
        return MTL::PixelFormatRGBA8Unorm;
    }
}

bool create_image(Device &device, ImageDesc &desc, Image &out) {
    MTL::TextureDescriptor *td = MTL::TextureDescriptor::texture2DDescriptor(
        to_pixel_format(desc.format), desc.width, desc.height, desc.mip_levels > 1
    );

    MTL::TextureUsage usage = MTL::TextureUsageShaderRead;
    if (rhi::Any(desc.usage, rhi::ImageUsage::ColorAttachment)) {
        usage |= MTL::TextureUsageRenderTarget;
    }
    if (rhi::Any(desc.usage, rhi::ImageUsage::DepthAttachment)) {
        usage |= MTL::TextureUsageRenderTarget;
    }
    if (rhi::Any(desc.usage, rhi::ImageUsage::Storage)) {
        usage |= MTL::TextureUsageShaderWrite;
    }
    td->setUsage(usage);
    td->setStorageMode(MTL::StorageModeShared);

    MTL::Texture *tex = device.handle->newTexture(td);
    td->release();
    if (!tex) {
        VEL_ERROR("Failed to create Metal texture ({}x{}x{})", desc.width, desc.height, desc.depth);
        return false;
    }

    out.desc = desc;
    out.handle = tex;
    if (!desc.name.empty()) {
        tex->setLabel(NS::String::string(desc.name.c_str(), NS::UTF8StringEncoding));
    }
    mt::reside(device, tex);
    return true;
}

void destroy_image(Device &device, Image &img) {
    destroy_image_views(device, img);
    if (img.handle) {
        mt::unreside(device, img.handle);
        img.handle->release();
        img.handle = nullptr;
    }
}

bool create_image_view(Device &device, const ImageViewDesc &desc, ImageView &out) {
    (void)device;
    out.desc = desc;
    out.handle = desc.image ? desc.image->handle : nullptr;
    if (out.handle) {
        out.handle->retain();
    }
    out.slot = mt::write_texture(device, out.handle, desc.type == rhi::ImageViewType::Storage);
    return out.handle != nullptr;
}

u64 handle_id(Device &device, ImageView &view) {
    if (view.slot == UINT64_MAX) {
        VEL_ERROR("handle_id: image view carries no heap slot (cached views are not shader-visible)");
        return 0;
    }
    return mt::heap_texture_id(device, view.slot);
}

void destroy_image_view(Device &device, ImageView &view) {
    mt::free_texture(device, view.slot);
    view.slot = UINT64_MAX;
    if (view.handle) {
        view.handle->release();
        view.handle = nullptr;
    }
}

static bool same_view_key(const ImageViewDesc &a, const ImageViewDesc &b) {
    return a.aspect == b.aspect && a.mip_start == b.mip_start && a.mip_count == b.mip_count &&
           a.layer_start == b.layer_start && a.layer_count == b.layer_count && a.dimension == b.dimension &&
           a.type == b.type;
}

ImageView get_cached_image_view(Device &device, Image &image, const ImageViewDesc &desc) {
    (void)device;
    ImageViewDesc key = desc;
    key.image = &image;
    for (auto &entry : image.view_cache) {
        if (same_view_key(entry.first, key)) {
            ImageView out{};
            out.desc = key;
            out.handle = entry.second;
            return out;
        }
    }
    if (image.handle) {
        image.handle->retain();
    }
    image.view_cache.emplace_back(key, image.handle);
    ImageView out{};
    out.desc = key;
    out.handle = image.handle;
    return out;
}

void destroy_image_views(Device &device, Image &image) {
    (void)device;
    for (auto &entry : image.view_cache) {
        if (entry.second) {
            entry.second->release();
        }
    }
    image.view_cache.clear();
}

void destroy_buffer_views(Device &device, Buffer &buffer) {
    (void)device;
    (void)buffer;
}

bool valid(const Image &img) {
    return img.handle != nullptr;
}

bool valid(const ImageView &view) {
    // cached views. no slot.
    return view.handle != nullptr;
}

} // namespace rhi

namespace mt {

void upload_image_data(rhi::Device &device, rhi::Image &image, const void *data, u32 bytes_per_row) {
    (void)device;
    if (!image.handle) {
        return;
    }
    MTL::Region region = MTL::Region(0, 0, image.desc.width, image.desc.height);
    image.handle->replaceRegion(region, 0, data, bytes_per_row);
}

} // namespace mt