blender-bridge / src / virtual_texture.cpp

#include "virtual_texture.h"

#include <godot_cpp/classes/image.hpp>
#include <godot_cpp/classes/rd_texture_format.hpp>
#include <godot_cpp/classes/rd_texture_view.hpp>
#include <godot_cpp/classes/resource_loader.hpp>
#include <godot_cpp/variant/utility_functions.hpp>

using namespace godot;

VirtualTexture::VirtualTexture() {
}

VirtualTexture::~VirtualTexture() {
	_free_gpu_resources();
}

// ── Bind methods ──────────────────────────────────────────────────────────────

void VirtualTexture::_bind_methods() {
	// Texture registration
	ClassDB::bind_method(D_METHOD("register_texture", "name", "width", "height"),
			&VirtualTexture::register_texture, DEFVAL(0), DEFVAL(0));
	ClassDB::bind_method(D_METHOD("unregister_texture", "vt_id"), &VirtualTexture::unregister_texture);
	ClassDB::bind_method(D_METHOD("get_texture_count"), &VirtualTexture::get_texture_count);

	// Invalidation
	ClassDB::bind_method(D_METHOD("invalidate_region", "texture_id", "tile_x", "tile_y", "w", "h", "mip"), &VirtualTexture::invalidate_region);
	ClassDB::bind_method(D_METHOD("invalidate_all", "texture_id"), &VirtualTexture::invalidate_all);

	// Atlas access
	ClassDB::bind_method(D_METHOD("get_atlas_texture"), &VirtualTexture::get_atlas_texture);
	ClassDB::bind_method(D_METHOD("get_page_table_texture"), &VirtualTexture::get_page_table_texture);

	// Material factory
	ClassDB::bind_method(D_METHOD("create_svt_material", "fallback_albedo", "roughness", "metallic"), &VirtualTexture::create_svt_material);

	// Feedback
	ClassDB::bind_method(D_METHOD("submit_feedback", "feedback_data"), &VirtualTexture::submit_feedback);

	// Tile upload
	ClassDB::bind_method(D_METHOD("upload_tile", "texture_id", "tile_x", "tile_y", "mip", "rgba_data"), &VirtualTexture::upload_tile);

	// CPU image reconstruction
	ClassDB::bind_method(D_METHOD("get_reconstructed_image", "vt_id"), &VirtualTexture::get_reconstructed_image);
	ClassDB::bind_method(D_METHOD("get_all_textures"), &VirtualTexture::get_all_textures);

	// Stats
	ClassDB::bind_method(D_METHOD("get_stats"), &VirtualTexture::get_stats);

	// Properties
	ClassDB::bind_method(D_METHOD("set_tile_size", "size"), &VirtualTexture::set_tile_size);
	ClassDB::bind_method(D_METHOD("get_tile_size"), &VirtualTexture::get_tile_size);
	ADD_PROPERTY(PropertyInfo(Variant::INT, "tile_size"), "set_tile_size", "get_tile_size");

	ClassDB::bind_method(D_METHOD("set_tile_budget", "budget"), &VirtualTexture::set_tile_budget);
	ClassDB::bind_method(D_METHOD("get_tile_budget"), &VirtualTexture::get_tile_budget);
	ADD_PROPERTY(PropertyInfo(Variant::INT, "tile_budget"), "set_tile_budget", "get_tile_budget");

	ClassDB::bind_method(D_METHOD("set_cache_grid_size", "size"), &VirtualTexture::set_cache_grid_size);
	ClassDB::bind_method(D_METHOD("get_cache_grid_size"), &VirtualTexture::get_cache_grid_size);
	ADD_PROPERTY(PropertyInfo(Variant::INT, "cache_grid_size"), "set_cache_grid_size", "get_cache_grid_size");

	ClassDB::bind_method(D_METHOD("set_mip_levels", "levels"), &VirtualTexture::set_mip_levels);
	ClassDB::bind_method(D_METHOD("get_mip_levels"), &VirtualTexture::get_mip_levels);
	ADD_PROPERTY(PropertyInfo(Variant::INT, "mip_levels"), "set_mip_levels", "get_mip_levels");

	// Signals
	ADD_SIGNAL(MethodInfo("tile_requested",
			PropertyInfo(Variant::INT, "texture_id"),
			PropertyInfo(Variant::INT, "tile_x"),
			PropertyInfo(Variant::INT, "tile_y"),
			PropertyInfo(Variant::INT, "mip")));
	ADD_SIGNAL(MethodInfo("tile_loaded",
			PropertyInfo(Variant::INT, "texture_id"),
			PropertyInfo(Variant::INT, "tile_x"),
			PropertyInfo(Variant::INT, "tile_y"),
			PropertyInfo(Variant::INT, "mip")));
}

// ── Properties ────────────────────────────────────────────────────────────────

void VirtualTexture::set_tile_size(int size) { _tile_size = size; }
int VirtualTexture::get_tile_size() const { return _tile_size; }
void VirtualTexture::set_tile_budget(int budget) { _tile_budget = budget; }
int VirtualTexture::get_tile_budget() const { return _tile_budget; }
void VirtualTexture::set_cache_grid_size(int size) { _cache_grid_size = size; }
int VirtualTexture::get_cache_grid_size() const { return _cache_grid_size; }
void VirtualTexture::set_mip_levels(int levels) { _mip_levels = levels; }
int VirtualTexture::get_mip_levels() const { return _mip_levels; }

// ── Texture registration ──────────────────────────────────────────────────────

int VirtualTexture::register_texture(const String &name, int width, int height) {
	// Dedup: return existing ID if already registered with same name
	for (auto &tex : _textures) {
		if (tex.active && tex.name == name) {
			UtilityFunctions::print("VirtualTexture: Reusing texture '", name, "' id=", tex.id);
			return tex.id;
		}
	}

	int id = _next_texture_id++;
	RegisteredTexture tex;
	tex.id = id;
	tex.name = name;
	tex.active = true;
	tex.width = width;
	tex.height = height;
	// Pre-allocate CPU image if dimensions known
	if (width > 0 && height > 0) {
		tex.cpu_image = Image::create(width, height, false, Image::FORMAT_RGBA8);
		tex.cpu_image->fill(Color(0, 0, 0, 0));
	}
	_textures.push_back(tex);
	UtilityFunctions::print("VirtualTexture: Registered texture '", name, "' as id=", id,
			" (", width, "x", height, ")");
	return id;
}

void VirtualTexture::unregister_texture(int vt_id) {
	for (auto &tex : _textures) {
		if (tex.id == vt_id) {
			tex.active = false;
			// Invalidate all tiles for this texture
			invalidate_all(vt_id);
			return;
		}
	}
}

int VirtualTexture::get_texture_count() const {
	int count = 0;
	for (const auto &tex : _textures) {
		if (tex.active)
			count++;
	}
	return count;
}

// ── Invalidation ──────────────────────────────────────────────────────────────

void VirtualTexture::invalidate_region(int texture_id, int tile_x, int tile_y, int w, int h, int mip) {
	// Evict from page table
	_page_table.invalidate_region(mip, tile_x, tile_y, w, h);

	// Add invalidation requests so tiles get re-fetched
	for (int ty = tile_y; ty < tile_y + h; ty++) {
		for (int tx = tile_x; tx < tile_x + w; tx++) {
			_request_queue.add_invalidated(texture_id, tx, ty, mip);
		}
	}
}

void VirtualTexture::invalidate_all(int texture_id) {
	_page_table.clear();
	// Free all cache slots for this texture
	// (simple approach — evict everything)
	_tile_cache.clear();
}

// ── Atlas / page table access ─────────────────────────────────────────────────

Ref<Texture2DRD> VirtualTexture::get_atlas_texture() const {
	return _atlas_tex;
}

Ref<Texture2DRD> VirtualTexture::get_page_table_texture() const {
	return _page_table_tex;
}

// ── Material factory ──────────────────────────────────────────────────────

Ref<ShaderMaterial> VirtualTexture::create_svt_material(const Color &fallback, float rough, float metal) {
	// Lazy-load the shader
	if (!_shader_loaded) {
		_svt_shader = ResourceLoader::get_singleton()->load("res://addons/blender_bridge/svt_lookup.gdshader");
		_shader_loaded = true;
	}
	if (_svt_shader.is_null()) {
		UtilityFunctions::push_warning("VirtualTexture: svt_lookup.gdshader not found at res://addons/blender_bridge/");
		return Ref<ShaderMaterial>();
	}

	Ref<ShaderMaterial> mat;
	mat.instantiate();
	mat->set_shader(_svt_shader);
	mat->set_shader_parameter("svt_atlas", _atlas_tex);
	mat->set_shader_parameter("svt_page_table", _page_table_tex);
	mat->set_shader_parameter("tile_size_uv", 1.0f / _cache_grid_size);
	mat->set_shader_parameter("page_table_size", (float)_page_table.tiles_at_mip(0));
	mat->set_shader_parameter("albedo_color", fallback);
	mat->set_shader_parameter("roughness", rough);
	mat->set_shader_parameter("metallic", metal);
	return mat;
}

// ── Tile upload ───────────────────────────────────────────────────────────────

Error VirtualTexture::upload_tile(int texture_id, int tile_x, int tile_y, int mip, const PackedByteArray &rgba_data) {
	int expected_size = _tile_size * _tile_size * 4; // RGBA8
	if (rgba_data.size() != expected_size) {
		UtilityFunctions::printerr("VirtualTexture: upload_tile got ",
				rgba_data.size(), " bytes, expected ", expected_size);
		return ERR_INVALID_PARAMETER;
	}

	if (!_gpu_initialized) {
		return ERR_UNCONFIGURED;
	}

	// Allocate or find existing cache slot
	TileId tile = { texture_id, tile_x, tile_y, mip };
	AllocResult alloc = _tile_cache.alloc(tile, _frame_count);
	if (!alloc.success) {
		return ERR_OUT_OF_MEMORY;
	}

	// Upload pixel data to atlas at the allocated location
	_upload_tile_to_atlas(alloc.location, rgba_data);

	// Update page table
	_page_table.map_with_fallback(mip, tile_x, tile_y, alloc.location);

	// Lock root mip tiles (top 2 mip levels never evicted)
	if (mip >= _mip_levels - 2) {
		_tile_cache.lock(alloc.location);
	}

	// Blit mip-0 tiles into the CPU reconstruction image
	if (mip == 0) {
		for (auto &tex : _textures) {
			if (tex.id == texture_id && tex.cpu_image.is_valid()) {
				Ref<Image> tile_img = Image::create_from_data(
						_tile_size, _tile_size, false, Image::FORMAT_RGBA8, rgba_data);
				int dst_x = tile_x * _tile_size;
				int dst_y = tile_y * _tile_size;
				int copy_w = MIN(_tile_size, tex.width - dst_x);
				int copy_h = MIN(_tile_size, tex.height - dst_y);
				if (copy_w > 0 && copy_h > 0) {
					tex.cpu_image->blit_rect(tile_img, Rect2i(0, 0, copy_w, copy_h),
							Vector2i(dst_x, dst_y));
				}
				break;
			}
		}
	}

	emit_signal("tile_loaded", texture_id, tile_x, tile_y, mip);
	_tiles_uploaded_this_frame++;
	return OK;
}

// ── Feedback submission ───────────────────────────────────────────────────────

void VirtualTexture::submit_feedback(const PackedByteArray &feedback_data) {
	if (feedback_data.size() < 4)
		return;

	int count = feedback_data.size() / 4;
	const uint32_t *data = reinterpret_cast<const uint32_t *>(feedback_data.ptr());
	_request_queue.add_from_feedback(data, count);
}

// ── Lifecycle ─────────────────────────────────────────────────────────────────

void VirtualTexture::_ready() {
	_page_table.init(_mip_levels);
	_tile_cache.init(_cache_grid_size, _cache_grid_size);
	_init_gpu_resources();
}

void VirtualTexture::_process(double delta) {
	_frame_count++;
	_tiles_uploaded_this_frame = 0;

	_process_tile_requests();

	// Upload updated page table to GPU each frame
	if (_gpu_initialized) {
		_upload_page_table_to_gpu();
	}
}

void VirtualTexture::_exit_tree() {
	_free_gpu_resources();
}

// ── GPU resource management ───────────────────────────────────────────────────

void VirtualTexture::_init_gpu_resources() {
	RenderingServer *rs = RenderingServer::get_singleton();
	if (!rs)
		return;

	RenderingDevice *rd = rs->get_rendering_device();
	if (!rd) {
		UtilityFunctions::printerr("VirtualTexture: No RenderingDevice available");
		return;
	}

	// ── Physical cache atlas ────────────────────────────────────────────
	// Size: cache_grid * tile_size in each dimension
	int atlas_size = _cache_grid_size * _tile_size;

	Ref<RDTextureFormat> atlas_fmt;
	atlas_fmt.instantiate();
	atlas_fmt->set_width(atlas_size);
	atlas_fmt->set_height(atlas_size);
	atlas_fmt->set_format(RenderingDevice::DATA_FORMAT_R8G8B8A8_UNORM);
	atlas_fmt->set_usage_bits(
			RenderingDevice::TEXTURE_USAGE_SAMPLING_BIT |
			RenderingDevice::TEXTURE_USAGE_CAN_UPDATE_BIT |
			RenderingDevice::TEXTURE_USAGE_CAN_COPY_TO_BIT);

	Ref<RDTextureView> atlas_view;
	atlas_view.instantiate();

	_atlas_rid = rd->texture_create(atlas_fmt, atlas_view);

	if (_atlas_rid.is_valid()) {
		_atlas_tex.instantiate();
		_atlas_tex->set_texture_rd_rid(_atlas_rid);
	}

	// ── Page table texture ──────────────────────────────────────────────
	// One texel per virtual tile at mip 0, RGBA8:
	// R = cache_x, G = cache_y, B = valid flag, A = 0
	int pt_size = _page_table.tiles_at_mip(0);

	Ref<RDTextureFormat> pt_fmt;
	pt_fmt.instantiate();
	pt_fmt->set_width(pt_size);
	pt_fmt->set_height(pt_size);
	pt_fmt->set_format(RenderingDevice::DATA_FORMAT_R8G8B8A8_UNORM);
	pt_fmt->set_usage_bits(
			RenderingDevice::TEXTURE_USAGE_SAMPLING_BIT |
			RenderingDevice::TEXTURE_USAGE_CAN_UPDATE_BIT);

	Ref<RDTextureView> pt_view;
	pt_view.instantiate();

	_page_table_rid = rd->texture_create(pt_fmt, pt_view);

	if (_page_table_rid.is_valid()) {
		_page_table_tex.instantiate();
		_page_table_tex->set_texture_rd_rid(_page_table_rid);
	}

	_gpu_initialized = _atlas_rid.is_valid() && _page_table_rid.is_valid();

	if (_gpu_initialized) {
		UtilityFunctions::print("VirtualTexture: GPU resources initialized (",
				atlas_size, "x", atlas_size, " atlas, ",
				pt_size, "x", pt_size, " page table)");
	} else {
		UtilityFunctions::printerr("VirtualTexture: Failed to create GPU textures");
	}
}

void VirtualTexture::_free_gpu_resources() {
	if (!_gpu_initialized)
		return;

	RenderingServer *rs = RenderingServer::get_singleton();
	if (!rs)
		return;

	RenderingDevice *rd = rs->get_rendering_device();
	if (!rd)
		return;

	if (_atlas_rid.is_valid()) {
		rd->free_rid(_atlas_rid);
		_atlas_rid = RID();
	}
	if (_page_table_rid.is_valid()) {
		rd->free_rid(_page_table_rid);
		_page_table_rid = RID();
	}

	_atlas_tex.unref();
	_page_table_tex.unref();
	_gpu_initialized = false;
}

void VirtualTexture::_upload_page_table_to_gpu() {
	RenderingServer *rs = RenderingServer::get_singleton();
	if (!rs)
		return;

	RenderingDevice *rd = rs->get_rendering_device();
	if (!rd)
		return;

	int pt_size = _page_table.tiles_at_mip(0);
	int data_size = pt_size * pt_size * 4; // RGBA8

	PackedByteArray pba;
	pba.resize(data_size);
	uint8_t *dst = pba.ptrw();

	_page_table.get_gpu_data(0, dst, data_size);

	// Upload to GPU
	rd->texture_update(_page_table_rid, 0, pba);
}

void VirtualTexture::_upload_tile_to_atlas(PhysicalLocation loc, const PackedByteArray &rgba_data) {
	RenderingServer *rs = RenderingServer::get_singleton();
	if (!rs)
		return;

	RenderingDevice *rd = rs->get_rendering_device();
	if (!rd)
		return;

	// Create a temporary source texture for this tile
	Ref<RDTextureFormat> tile_fmt;
	tile_fmt.instantiate();
	tile_fmt->set_width(_tile_size);
	tile_fmt->set_height(_tile_size);
	tile_fmt->set_format(RenderingDevice::DATA_FORMAT_R8G8B8A8_UNORM);
	tile_fmt->set_usage_bits(
			RenderingDevice::TEXTURE_USAGE_CAN_COPY_FROM_BIT |
			RenderingDevice::TEXTURE_USAGE_CAN_UPDATE_BIT);

	Ref<RDTextureView> tile_view;
	tile_view.instantiate();

	TypedArray<PackedByteArray> data_layers;
	data_layers.push_back(rgba_data);

	RID tile_rid = rd->texture_create(tile_fmt, tile_view, data_layers);
	if (!tile_rid.is_valid()) {
		UtilityFunctions::printerr("VirtualTexture: Failed to create temp tile texture");
		return;
	}

	// Copy tile into atlas at the physical location
	// Atlas position = loc * tile_size
	int dst_x = loc.x * _tile_size;
	int dst_y = loc.y * _tile_size;

	rd->texture_copy(
			tile_rid, _atlas_rid,
			Vector3(0, 0, 0), // src offset
			Vector3(dst_x, dst_y, 0), // dst offset
			Vector3(_tile_size, _tile_size, 1), // size
			0, 0, // src/dst mip
			0, 0 // src/dst layer
	);

	// Free temporary texture
	rd->free_rid(tile_rid);
}

// ── Per-frame pipeline ────────────────────────────────────────────────────────

void VirtualTexture::_process_tile_requests() {
	// Get prioritized batch within budget
	auto batch = _request_queue.get_next_batch(_tile_budget);

	for (const auto &req : batch) {
		// Allocate a cache slot
		AllocResult alloc = _tile_cache.alloc(req.tile, _frame_count);
		if (!alloc.success)
			continue;

		// If eviction happened, update page table for evicted tile
		if (alloc.was_eviction) {
			// The evicted tile's page table entry will be stale — it'll
			// get a fallback or be re-requested via feedback next frame
		}

		// Map in page table
		_page_table.map_with_fallback(
				req.tile.mip, req.tile.tile_x, req.tile.tile_y,
				alloc.location);

		// Signal that we need this tile's data
		emit_signal("tile_requested",
				req.tile.texture_id, req.tile.tile_x,
				req.tile.tile_y, req.tile.mip);

		_tiles_uploaded_this_frame++;
	}
}

// ── CPU image reconstruction ──────────────────────────────────────────────────

Ref<Image> VirtualTexture::get_reconstructed_image(int vt_id) const {
	for (const auto &tex : _textures) {
		if (tex.id == vt_id && tex.active) {
			return tex.cpu_image;
		}
	}
	return Ref<Image>();
}

TypedArray<Dictionary> VirtualTexture::get_all_textures() const {
	TypedArray<Dictionary> result;
	for (const auto &tex : _textures) {
		if (!tex.active)
			continue;
		Dictionary d;
		d["id"] = tex.id;
		d["name"] = tex.name;
		d["width"] = tex.width;
		d["height"] = tex.height;
		d["has_image"] = tex.cpu_image.is_valid();
		result.push_back(d);
	}
	return result;
}

Ref<Shader> VirtualTexture::get_svt_shader() const {
	return _svt_shader;
}

// ── Stats ─────────────────────────────────────────────────────────────────────

Dictionary VirtualTexture::get_stats() const {
	Dictionary stats;
	stats["frame"] = (int)_frame_count;
	stats["tiles_uploaded_this_frame"] = _tiles_uploaded_this_frame;
	stats["cache_occupied"] = _tile_cache.occupied_count();
	stats["cache_total"] = _tile_cache.total_slots();
	stats["cache_locked"] = _tile_cache.locked_count();
	stats["pending_requests"] = _request_queue.pending_count();
	stats["registered_textures"] = get_texture_count();
	return stats;
}