#include "stb_image_loader.h"
#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#include "interface/image_interface.h"

sg_pixel_format stb_image_info::get_pixel_format() const {
	// 根据通道数和是否HDR选择合适的像素格式
	if (is_hdr) {
		switch (channels) {
			case 1:
				return SG_PIXELFORMAT_R32F;
			case 2:
				return SG_PIXELFORMAT_RG32F;
			case 3:
			case 4:
				return SG_PIXELFORMAT_RGBA32F;
			default:
				return SG_PIXELFORMAT_RGBA32F;
		}
	}

	// 16位图像
	if (is_16_bit) {
		switch (channels) {
			case 1:
				return SG_PIXELFORMAT_R16;
			case 2:
				return SG_PIXELFORMAT_RG16;
			case 3:
			case 4:
				return SG_PIXELFORMAT_RGBA16;
			default:
				return SG_PIXELFORMAT_RGBA16;
		}
	}

	// 8位图像
	switch (channels) {
		case 1:
			return SG_PIXELFORMAT_R8;
		case 2:
			return SG_PIXELFORMAT_RG8;
		case 3:
		case 4:
			return SG_PIXELFORMAT_RGBA8;
		default:
			return SG_PIXELFORMAT_RGBA8;
	}
}

void stb_image_loader::init(void* in_data, int32_t in_size) {
	data_ = static_cast<stbi_uc*>(in_data);
	size_ = in_size;
}

stb_image_info stb_image_loader::get_info() const {
	stb_image_info info{};
	stbi_info_from_memory(data_, size_, &info.width, &info.height, &info.channels);
	info.is_hdr = is_hdr();
	info.is_16_bit = is_16_bit();
	return info;
}

bool stb_image_loader::is_hdr() const {
	return stbi_is_hdr_from_memory(data_, size_);
}

bool stb_image_loader::is_16_bit() const {
	return stbi_is_16_bit_from_memory(data_, size_);
}

float* stb_image_loader::load_hdr() const {
	int width, height, channels;
	return stbi_loadf_from_memory(data_, size_, &width, &height, &channels, 4);
}

uint16_t* stb_image_loader::load_16_bit() const {
	int width, height, channels;
	return stbi_load_16_from_memory(data_, size_, &width, &height, &channels, 4);
}

uint8_t* stb_image_loader::load_8_bit() const {
	int width, height, channels;
	return stbi_load_from_memory(data_, size_, &width, &height, &channels, 4);
}

struct stb_heap_image {
	stb_image_info info;
	void* data;

	~stb_heap_image() {
		if (data)
			stbi_image_free(data);
	}
};

void delete_stb_image(image_heap_t* in_data) {
	if (in_data->data)
		stbi_image_free(in_data->data);
	delete in_data;
}

std::shared_ptr<image_heap_t> stb_image_loader::load() const {
	std::shared_ptr<image_heap_t> heap(new image_heap_t, delete_stb_image);
	const auto& info = get_info();
	heap->width = info.width;
	heap->height = info.height;
	heap->pixel_format = info.get_pixel_format();

	if (is_hdr())
		heap->data = load_hdr();
	else if (is_16_bit())
		heap->data = load_16_bit();
	else
		heap->data = load_8_bit();
	return heap;
}

std::shared_ptr<image_heap_t> load_image(const std::span<std::byte>& in_raw_data) {
	stb_image_loader loader;
	loader.init(in_raw_data.data(), static_cast<int32_t>(in_raw_data.size()));
	return loader.load();
}
