#include <psemek/wgpu/device.hpp>
#include <psemek/wgpu/external/webgpu.h>

#include <cstring>

namespace psemek::wgpu
{

	queue device::get_queue()
	{
		return queue(wgpuDeviceGetQueue((WGPUDevice)get()));
	}

	bind_group device::create_bind_group(bind_group::descriptor const & desc)
	{
		std::vector<WGPUBindGroupEntry> entries;
		for (auto const & entry_in : desc.entries)
		{
			auto & entry_out = entries.emplace_back();
			entry_out.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(entry_in.chain);
			entry_out.binding = entry_in.binding;
			entry_out.buffer = (WGPUBuffer)entry_in.buffer.get();
			entry_out.offset = entry_in.offset;
			entry_out.size = entry_in.size;
			entry_out.sampler = (WGPUSampler)entry_in.sampler.get();
			entry_out.textureView = (WGPUTextureView)entry_in.texture_view.get();
		}

		WGPUBindGroupDescriptor descriptor = {};
		descriptor.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.chain);
		descriptor.label = desc.label.data();
		descriptor.layout = (WGPUBindGroupLayout)desc.layout.get();
		descriptor.entryCount = entries.size();
		descriptor.entries = entries.data();

		return bind_group(wgpuDeviceCreateBindGroup((WGPUDevice)get(), &descriptor));
	}

	bind_group_layout device::create_bind_group_layout(bind_group_layout::descriptor const & desc)
	{
		std::vector<WGPUBindGroupLayoutEntry> entries;
		for (auto const & entry_in : desc.entries)
		{
			auto & entry_out = entries.emplace_back();
			entry_out.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(entry_in.chain);
			entry_out.binding = entry_in.binding;
			entry_out.visibility = (WGPUShaderStageFlags)entry_in.visibility;
			entry_out.buffer.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(entry_in.buffer.chain);
			entry_out.buffer.type = (WGPUBufferBindingType)entry_in.buffer.type;
			entry_out.buffer.hasDynamicOffset = entry_in.buffer.has_dynamic_offset;
			entry_out.buffer.minBindingSize = entry_in.buffer.min_binding_size;
			entry_out.sampler.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(entry_in.sampler.chain);
			entry_out.sampler.type = (WGPUSamplerBindingType)entry_in.sampler.type;
			entry_out.texture.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(entry_in.texture.chain);
			entry_out.texture.sampleType = (WGPUTextureSampleType)entry_in.texture.sample_type;
			entry_out.texture.viewDimension = (WGPUTextureViewDimension)entry_in.texture.view_dimension;
			entry_out.texture.multisampled = entry_in.texture.multisampled;
			entry_out.storageTexture.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(entry_in.storage_texture.chain);
			entry_out.storageTexture.access = (WGPUStorageTextureAccess)entry_in.storage_texture.access;
			entry_out.storageTexture.format = (WGPUTextureFormat)entry_in.storage_texture.format;
			entry_out.storageTexture.viewDimension = (WGPUTextureViewDimension)entry_in.storage_texture.view_dimension;
		}

		WGPUBindGroupLayoutDescriptor descriptor = {};
		descriptor.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.chain);
		descriptor.label = desc.label.data();
		descriptor.entryCount = entries.size();
		descriptor.entries = entries.data();

		return bind_group_layout(wgpuDeviceCreateBindGroupLayout((WGPUDevice)get(), &descriptor));
	}

	buffer device::create_buffer(buffer::descriptor const & desc)
	{
		WGPUBufferDescriptor descriptor = {};
		descriptor.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.chain);
		descriptor.label = desc.label.data();
		descriptor.usage = (WGPUBufferUsageFlags)desc.usage;
		descriptor.size = desc.size;
		descriptor.mappedAtCreation = desc.mapped_at_creation;
		return buffer(wgpuDeviceCreateBuffer((WGPUDevice)get(), &descriptor));
	}

	command_encoder device::create_command_encoder(command_encoder::descriptor const & desc)
	{
		WGPUCommandEncoderDescriptor descriptor = {};
		descriptor.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.chain);
		descriptor.label = desc.label.data();
		return command_encoder(wgpuDeviceCreateCommandEncoder((WGPUDevice)get(), &descriptor));
	}

	namespace
	{

		void fill_compute_pipeline_descriptor(compute_pipeline::descriptor const & desc, std::vector<WGPUConstantEntry> & constants, WGPUComputePipelineDescriptor & descriptor)
		{
			for (auto const & constant_in : desc.compute.constants)
			{
				auto & constant_out = constants.emplace_back();
				constant_out.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(constant_in.chain);
				constant_out.key = constant_in.key.data();
				constant_out.value = constant_in.value;
			}

			descriptor.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.chain);
			descriptor.label = desc.label.data();
			descriptor.layout = (WGPUPipelineLayout)desc.layout.get();
			descriptor.compute.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.compute.chain);
			descriptor.compute.module = (WGPUShaderModule)desc.compute.module.get();
			descriptor.compute.entryPoint = desc.compute.entry_point.data();
			descriptor.compute.constantCount = constants.size();
			descriptor.compute.constants = constants.data();
		}

	}

	compute_pipeline device::create_compute_pipeline(compute_pipeline::descriptor const & desc)
	{
		std::vector<WGPUConstantEntry> constants;
		WGPUComputePipelineDescriptor descriptor = {};
		fill_compute_pipeline_descriptor(desc, constants, descriptor);

		return compute_pipeline(wgpuDeviceCreateComputePipeline((WGPUDevice)get(), &descriptor));
	}

	void device::create_compute_pipeline_async(compute_pipeline::descriptor const & desc, create_compute_pipeline_async_callback const & callback)
	{
		std::vector<WGPUConstantEntry> constants;
		WGPUComputePipelineDescriptor descriptor = {};
		fill_compute_pipeline_descriptor(desc, constants, descriptor);

		auto userdata = new create_compute_pipeline_async_callback(callback);

		auto real_callback = [](WGPUCreatePipelineAsyncStatus status, WGPUComputePipeline pipeline, char const * message, void * userdata)
		{
			std::unique_ptr<create_compute_pipeline_async_callback> callback((create_compute_pipeline_async_callback *)userdata);
			if (*callback) (*callback)((create_pipeline_async_status)status, compute_pipeline(pipeline), std::string(message ? message : ""));
		};

		wgpuDeviceCreateComputePipelineAsync((WGPUDevice)get(), &descriptor, real_callback, userdata);
	}

	pipeline_layout device::create_pipeline_layout(pipeline_layout::descriptor const & desc)
	{
		WGPUPipelineLayoutDescriptor descriptor = {};
		descriptor.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.chain);
		descriptor.label = desc.label.data();
		descriptor.bindGroupLayoutCount = desc.layouts.size();
		static_assert(sizeof(WGPUBindGroupLayout) == sizeof(bind_group_layout));
		descriptor.bindGroupLayouts = (WGPUBindGroupLayout const *)desc.layouts.data();
		return pipeline_layout(wgpuDeviceCreatePipelineLayout((WGPUDevice)get(), &descriptor));
	}

	query_set device::create_query_set(query_set::descriptor const & desc)
	{
		WGPUQuerySetDescriptor descriptor = {};
		descriptor.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.chain);
		descriptor.label = desc.label.data();
		descriptor.type = (WGPUQueryType)desc.type;
		descriptor.count = desc.count;
		return query_set(wgpuDeviceCreateQuerySet((WGPUDevice)get(), &descriptor));
	}

	namespace
	{

		void fill_render_pipeline_descriptor(render_pipeline::descriptor const & desc,
			WGPURenderPipelineDescriptor & descriptor,
			std::vector<WGPUConstantEntry> & vertex_constants,
			std::vector<WGPUVertexBufferLayout> & vertex_buffers,
			WGPUDepthStencilState & depth_stencil_state,
			WGPUFragmentState & fragment_state,
			std::vector<WGPUConstantEntry> & fragment_constants,
			std::vector<WGPUColorTargetState> & color_targets)
		{
			descriptor.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.chain);
			descriptor.label = desc.label.data();
			descriptor.layout = (WGPUPipelineLayout)desc.layout.get();

			for (auto const & constant_in : desc.vertex.constants)
			{
				auto & constant_out = vertex_constants.emplace_back();
				constant_out.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(constant_in.chain);
				constant_out.key = constant_in.key.data();
				constant_out.value = constant_in.value;
			}

			for (auto const & buffer_in : desc.vertex.buffers)
			{
				auto & buffer_out = vertex_buffers.emplace_back();
				buffer_out.arrayStride = buffer_in.array_stride;
				buffer_out.stepMode = (WGPUVertexStepMode)buffer_in.step_mode;
				buffer_out.attributeCount = buffer_in.attributes.size();
				static_assert(sizeof(WGPUVertexAttribute) == sizeof(vertex_attribute));
				buffer_out.attributes = (WGPUVertexAttribute const *)buffer_in.attributes.data();
			}

			descriptor.vertex.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.vertex.chain);
			descriptor.vertex.module = (WGPUShaderModule)desc.vertex.module.get();
			descriptor.vertex.entryPoint = desc.vertex.entry_point.data();
			descriptor.vertex.constantCount = vertex_constants.size();
			descriptor.vertex.constants = vertex_constants.data();
			descriptor.vertex.bufferCount = vertex_buffers.size();
			descriptor.vertex.buffers = vertex_buffers.data();

			descriptor.primitive.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.primitive.chain);
			descriptor.primitive.topology = (WGPUPrimitiveTopology)desc.primitive.topology;
			descriptor.primitive.stripIndexFormat = (WGPUIndexFormat)desc.primitive.strip_index_format;
			descriptor.primitive.frontFace = (WGPUFrontFace)desc.primitive.front_face;
			descriptor.primitive.cullMode = (WGPUCullMode)desc.primitive.cull_mode;

			if (desc.depth_stencil)
			{
				depth_stencil_state.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.depth_stencil->chain);
				depth_stencil_state.format = (WGPUTextureFormat)desc.depth_stencil->format;
				depth_stencil_state.depthWriteEnabled = desc.depth_stencil->depth_write;
				depth_stencil_state.depthCompare = (WGPUCompareFunction)desc.depth_stencil->depth_compare;
				depth_stencil_state.stencilFront.compare = (WGPUCompareFunction)desc.depth_stencil->stencil_front.compare;
				depth_stencil_state.stencilFront.failOp = (WGPUStencilOperation)desc.depth_stencil->stencil_front.fail_op;
				depth_stencil_state.stencilFront.depthFailOp = (WGPUStencilOperation)desc.depth_stencil->stencil_front.depth_fail_op;
				depth_stencil_state.stencilFront.passOp = (WGPUStencilOperation)desc.depth_stencil->stencil_front.pass_op;
				depth_stencil_state.stencilBack.compare = (WGPUCompareFunction)desc.depth_stencil->stencil_back.compare;
				depth_stencil_state.stencilBack.failOp = (WGPUStencilOperation)desc.depth_stencil->stencil_back.fail_op;
				depth_stencil_state.stencilBack.depthFailOp = (WGPUStencilOperation)desc.depth_stencil->stencil_back.depth_fail_op;
				depth_stencil_state.stencilBack.passOp = (WGPUStencilOperation)desc.depth_stencil->stencil_back.pass_op;
				depth_stencil_state.stencilReadMask = desc.depth_stencil->stencil_read_mask;
				depth_stencil_state.stencilWriteMask = desc.depth_stencil->stencil_write_mask;
				depth_stencil_state.depthBias = desc.depth_stencil->depth_bias;
				depth_stencil_state.depthBiasSlopeScale = desc.depth_stencil->depth_bias_slope_scale;
				depth_stencil_state.depthBiasClamp = desc.depth_stencil->depth_bias_clamp;

				descriptor.depthStencil = &depth_stencil_state;
			}

			descriptor.multisample.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.multisample.chain);
			descriptor.multisample.count = desc.multisample.count;
			descriptor.multisample.mask = desc.multisample.mask;
			descriptor.multisample.alphaToCoverageEnabled = desc.multisample.alpha_to_coverage;

			if (desc.fragment)
			{
				for (auto const & constant_in : desc.fragment->constants)
				{
					auto & constant_out = fragment_constants.emplace_back();
					constant_out.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(constant_in.chain);
					constant_out.key = constant_in.key.data();
					constant_out.value = constant_in.value;
				}

				for (auto const & target_in : desc.fragment->targets)
				{
					auto & target_out = color_targets.emplace_back();
					target_out.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(target_in.chain);
					target_out.format = (WGPUTextureFormat)target_in.format;
					static_assert(sizeof(WGPUBlendState) == sizeof(blend_state));
					target_out.blend = (WGPUBlendState *)(target_in.blend ? &(*target_in.blend) : nullptr);
					target_out.writeMask = (WGPUColorWriteMaskFlags)target_in.write_mask;
				}

				fragment_state.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.fragment->chain);
				fragment_state.module = (WGPUShaderModule)desc.fragment->module.get();
				fragment_state.entryPoint = desc.fragment->entry_point.data();
				fragment_state.constantCount = fragment_constants.size();
				fragment_state.constants = fragment_constants.data();
				fragment_state.targetCount = color_targets.size();
				fragment_state.targets = color_targets.data();

				descriptor.fragment = &fragment_state;
			}
		}

	}

	render_pipeline device::create_render_pipeline(render_pipeline::descriptor const & desc)
	{
		WGPURenderPipelineDescriptor descriptor = {};
		std::vector<WGPUConstantEntry> vertex_constants;
		std::vector<WGPUVertexBufferLayout> vertex_buffers;
		WGPUDepthStencilState depth_stencil_state = {};
		WGPUFragmentState fragment_state = {};
		std::vector<WGPUConstantEntry> fragment_constants;
		std::vector<WGPUColorTargetState> color_targets;

		fill_render_pipeline_descriptor(desc, descriptor, vertex_constants, vertex_buffers, depth_stencil_state, fragment_state, fragment_constants, color_targets);

		return render_pipeline(wgpuDeviceCreateRenderPipeline((WGPUDevice)get(), &descriptor));
	}

	void device::create_render_pipeline_async(render_pipeline::descriptor const & desc, create_render_pipeline_async_callback const & callback)
	{
		WGPURenderPipelineDescriptor descriptor = {};
		std::vector<WGPUConstantEntry> vertex_constants;
		std::vector<WGPUVertexBufferLayout> vertex_buffers;
		WGPUDepthStencilState depth_stencil_state = {};
		WGPUFragmentState fragment_state = {};
		std::vector<WGPUConstantEntry> fragment_constants;
		std::vector<WGPUColorTargetState> color_targets;

		fill_render_pipeline_descriptor(desc, descriptor, vertex_constants, vertex_buffers, depth_stencil_state, fragment_state, fragment_constants, color_targets);

		auto userdata = new create_render_pipeline_async_callback(callback);

		auto real_callback = [](WGPUCreatePipelineAsyncStatus status, WGPURenderPipeline pipeline, char const * message, void * userdata)
		{
			std::unique_ptr<create_render_pipeline_async_callback> callback((create_render_pipeline_async_callback *)userdata);
			if (*callback) (*callback)((create_pipeline_async_status)status, render_pipeline(pipeline), std::string(message ? message : ""));
		};

		wgpuDeviceCreateRenderPipelineAsync((WGPUDevice)get(), &descriptor, real_callback, userdata);
	}

	render_bundle_encoder device::create_render_bundle_encoder(render_bundle_encoder::descriptor const & desc)
	{
		WGPURenderBundleEncoderDescriptor descriptor = {};
		descriptor.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.chain);
		descriptor.label = desc.label.data();
		descriptor.colorFormatCount = desc.color_formats.size();
		descriptor.colorFormats = (WGPUTextureFormat const *)desc.color_formats.data();
		descriptor.depthStencilFormat = (WGPUTextureFormat)desc.depth_stencil_format;
		descriptor.sampleCount = desc.sample_count;
		descriptor.depthReadOnly = desc.depth_read_only;
		descriptor.stencilReadOnly = desc.stencil_read_only;
		return render_bundle_encoder(wgpuDeviceCreateRenderBundleEncoder((WGPUDevice)get(), &descriptor));
	}

	sampler device::create_sampler(sampler::descriptor const & desc)
	{
		WGPUSamplerDescriptor descriptor = {};
		descriptor.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.chain);
		descriptor.label = desc.label.data();
		descriptor.addressModeU = (WGPUAddressMode)desc.address_mode_u;
		descriptor.addressModeV = (WGPUAddressMode)desc.address_mode_v;
		descriptor.addressModeW = (WGPUAddressMode)desc.address_mode_w;
		descriptor.magFilter = (WGPUFilterMode)desc.mag_filter;
		descriptor.minFilter = (WGPUFilterMode)desc.min_filter;
		descriptor.mipmapFilter = (WGPUMipmapFilterMode)desc.mipmap_filter;
		descriptor.lodMinClamp = desc.lod_clamp.min;
		descriptor.lodMaxClamp = desc.lod_clamp.max;
		descriptor.compare = (WGPUCompareFunction)desc.compare;
		descriptor.maxAnisotropy = desc.maxAnisotropy;
		return sampler(wgpuDeviceCreateSampler((WGPUDevice)get(), &descriptor));
	}

	shader_module device::create_shader_module(shader_module::descriptor const & desc)
	{
		std::vector<WGPUShaderModuleCompilationHint> hints;
		for (auto const & hint_in : desc.hints)
		{
			auto & hint = hints.emplace_back();
			hint.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(hint_in.chain);
			hint.entryPoint = hint_in.entry_point.data();
			hint.layout = (WGPUPipelineLayout)hint_in.layout.get();
		}

		WGPUShaderModuleDescriptor descriptor = {};
		descriptor.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.chain);
		descriptor.label = desc.label.data();
		descriptor.hintCount = hints.size();
		descriptor.hints = hints.data();
		return shader_module(wgpuDeviceCreateShaderModule((WGPUDevice)get(), &descriptor));
	}

	texture device::create_texture(texture::descriptor const & desc)
	{
		WGPUTextureDescriptor descriptor = {};
		descriptor.nextInChain = (WGPUChainedStruct const *)detail::fill_chain(desc.chain);
		descriptor.label = desc.label.data();
		descriptor.usage = (WGPUTextureUsageFlags)desc.usage;
		descriptor.dimension = (WGPUTextureDimension)desc.dimension;
		descriptor.size = {desc.size[0], desc.size[1], desc.size[2]};
		descriptor.format = (WGPUTextureFormat)desc.format;
		descriptor.mipLevelCount = desc.mip_level_count;
		descriptor.sampleCount = desc.sample_count;
		descriptor.viewFormatCount = desc.view_formats.size();
		descriptor.viewFormats = (WGPUTextureFormat *)desc.view_formats.data();
		return texture(wgpuDeviceCreateTexture((WGPUDevice)get(), &descriptor));
	}

	void device::destroy()
	{
		wgpuDeviceDestroy((WGPUDevice)get());
	}

	std::vector<feature> device::enumerate_features()
	{
		std::size_t count = wgpuDeviceEnumerateFeatures((WGPUDevice)get(), nullptr);
		std::vector<feature> result(count);
		wgpuDeviceEnumerateFeatures((WGPUDevice)get(), (WGPUFeatureName *)result.data());
		return result;
	}

	limits device::get_limits()
	{
		WGPUSupportedLimits limits = {};
		wgpuDeviceGetLimits((WGPUDevice)get(), &limits);

		// TODO: support out chain
		wgpu::limits result;
		static_assert(sizeof(result) == sizeof(limits.limits));
		std::memcpy((char *)&result, (char *)&limits.limits, sizeof(result));
		return result;
	}

	bool device::has_feature(feature feature)
	{
		return wgpuDeviceHasFeature((WGPUDevice)get(), (WGPUFeatureName)feature);
	}

	void device::push_error_scope(error_filter filter)
	{
		wgpuDevicePushErrorScope((WGPUDevice)get(), (WGPUErrorFilter)filter);
	}

	void device::pop_error_scope(error_callback const & callback)
	{
		auto userdata = new error_callback(callback);

		auto real_callback = [](WGPUErrorType type, char const * message, void * userdata)
		{
			std::unique_ptr<error_callback> callback((error_callback *)userdata);
			if (*callback) (*callback)((error_type)type, message ? message : "");
		};

		wgpuDevicePopErrorScope((WGPUDevice)get(), real_callback, userdata);
	}

	void device::set_uncaptured_error_callback(error_callback const & callback)
	{
		auto userdata = new error_callback(callback);

		auto real_callback = [](WGPUErrorType type, char const * message, void * userdata)
		{
			std::unique_ptr<error_callback> callback((error_callback *)userdata);
			if (*callback) (*callback)((error_type)type, message ? message : "");
		};

		wgpuDeviceSetUncapturedErrorCallback((WGPUDevice)get(), real_callback, userdata);
	}

	void device::set_label(std::string const & label)
	{
		wgpuDeviceSetLabel((WGPUDevice)get(), label.data());
	}

	void device::reference(void * ptr)
	{
		wgpuDeviceReference((WGPUDevice)ptr);
	}

	void device::release(void * ptr)
	{
		wgpuDeviceRelease((WGPUDevice)ptr);
	}

}