#pragma once

#include <cstdint>

#include "ShaderCommon.h"
#include "Common/Math/math_util.h"
#include "GPU/GPUState.h"

// Used by the "modern" backends that use uniform buffers. They can share this without issue.

enum : uint64_t {
	DIRTY_BASE_UNIFORMS =
	DIRTY_WORLDMATRIX | DIRTY_PROJTHROUGHMATRIX | DIRTY_VIEWMATRIX | DIRTY_TEXMATRIX | DIRTY_ALPHACOLORREF |
	DIRTY_PROJMATRIX | DIRTY_FOGCOLOR | DIRTY_FOGCOEF | DIRTY_TEXENV | DIRTY_TEX_ALPHA_MUL | DIRTY_STENCILREPLACEVALUE |
	DIRTY_ALPHACOLORMASK | DIRTY_SHADERBLEND | DIRTY_COLORWRITEMASK | DIRTY_UVSCALEOFFSET | DIRTY_TEXCLAMP | DIRTY_MATAMBIENTALPHA |
	DIRTY_DEPAL | DIRTY_VIEWPORT_UNIFORMS | DIRTY_RASTER_OFFSET,
	DIRTY_LIGHT_UNIFORMS =
	DIRTY_LIGHT_CONTROL | DIRTY_LIGHT0 | DIRTY_LIGHT1 | DIRTY_LIGHT2 | DIRTY_LIGHT3 |
	DIRTY_MATDIFFUSE | DIRTY_MATSPECULAR | DIRTY_MATEMISSIVE | DIRTY_AMBIENT,
};

// Currently 480 bytes.
// Every line here is a 4-float.
struct alignas(16) UB_VS_FS_Base {
	float proj[16];
	float view[12];
	float world[12];
	float tex[12];
	float xywh[4];  // later, we could invert w and h here to avoid division.
	float vpScale[3]; float NaN;
	float vpOffset[3]; float padding2;
	float rasterOffset[2]; float minZmaxZ[2];
	float uvScaleOffset[4];
	float matAmbient[4];
	uint32_t padding3; uint32_t depal_mask_shift_off_fmt;  // 4 params packed into one.
	uint32_t colorWriteMask; float mipBias;
	// Fragment data
	float texNoAlpha; float texMul; float padding4[2];  // this vec4 will hold ubershader stuff. We won't use integer flags in the fragment shader.
	float fogColor[3]; uint32_t alphaColorRef;
	float texEnvColor[3]; uint32_t colorTestMask;
	float texClamp[4];
	float texClampOffset[2]; float fogCoef[2];
	float blendFixA[3]; float stencilReplaceValue;
	float blendFixB[3]; float rotation;
	// VR stuff is to go here, later. For normal drawing, we can then get away
	// with just uploading the first X bytes of the struct (up to and including fogCoef).
};
static_assert(sizeof(UB_VS_FS_Base) <= 432, "UB_VS_FS_Base should be 432 bytes");

static const char * const ub_baseStr =
R"(  mat4 u_proj;
  mat3x4 u_view;
  mat3x4 u_world;
  mat3x4 u_texmtx;
  vec4 u_xywh;
  vec3 u_vpScale; float u_NaN; // w = offsetX
  vec4 u_vpOffset; // w = offsetY
  vec2 u_rasterOffset; vec2 u_minZmaxZ;
  vec4 u_uvscaleoffset;
  vec4 u_matambientalpha;
  uint u_spline_counts;
  uint u_depal_mask_shift_off_fmt;
  uint u_colorWriteMask;
  float u_mipBias;
  vec2 u_texNoAlphaMul; float pad1; float pad2;
  vec3 u_fogcolor;  uint u_alphacolorref;
  vec3 u_texenv;    uint u_alphacolormask;
  vec4 u_texclamp;
  vec2 u_texclampoff; vec2 u_fogcoef;
  vec3 u_blendFixA; float u_stencilReplaceValue;
  vec3 u_blendFixB; float u_rotation;
)";

// 512 bytes. Would like to shrink more. Some colors only have 8-bit precision and we expand
// them to float unnecessarily, could just as well expand in the shader.
struct alignas(16) UB_VS_Lights {
	float ambientColor[4];
	float materialDiffuse[4];
	float materialSpecular[4];
	float materialEmissive[3];
	uint32_t lightControl;
	float lpos[4][4];
	float ldir[4][4];
	float latt[4][4];
	float lightAngle_SpotCoef[4][4];   // TODO: Merge with lightSpotCoef, use .xy
	float lightAmbient[4][4];
	float lightDiffuse[4][4];
	float lightSpecular[4][4];
};
static_assert(sizeof(UB_VS_Lights) == 512);  // it's ok to optimize this, it's just an assumption check.

static const char * const ub_vs_lightsStr =
R"(  vec4 u_ambient;
  vec3 u_matdiffuse;
  vec4 u_matspecular;
  vec3 u_matemissive;
  uint u_lightControl;  // light ubershader control bits
  vec3 u_lightpos[4];
  vec3 u_lightdir[4];
  vec3 u_lightatt[4];
  vec4 u_lightangle_spotCoef[4];
  vec3 u_lightambient[4];
  vec3 u_lightdiffuse[4];
  vec3 u_lightspecular[4];
)";

// With some cleverness, we could get away with uploading just half this when only the four or five first
// bones are being used. This is 384b.
struct alignas(16) UB_VS_Bones {
	float bones[8][12];
};
static_assert(sizeof(UB_VS_Bones) == 384);  // No way to optimize this further.

static const char * const ub_vs_bonesStr =
R"(	mat3x4 u_bone0; mat3x4 u_bone1; mat3x4 u_bone2; mat3x4 u_bone3; mat3x4 u_bone4; mat3x4 u_bone5; mat3x4 u_bone6; mat3x4 u_bone7; mat3x4 u_bone8;
)";

// useBufferedRendering is only used to determine the rotation uniform.
void BaseUpdateUniforms(UB_VS_FS_Base *ub, uint64_t dirtyUniforms, bool useBufferedRendering);
void LightUpdateUniforms(UB_VS_Lights *ub, uint64_t dirtyUniforms);
void BoneUpdateUniforms(UB_VS_Bones *ub, uint64_t dirtyUniforms);

uint32_t PackLightControlBits();
uint32_t PackDepalBits();

void UpdateFogCoef(const GEState &state, float fogCoef[2]);

// This happens so much that I want it inline.
inline void UpdateUVScaleOff(const GEState &state, float uvScaleOff[4]) {
	float widthFactor;
	float heightFactor;
	if (gstate_c.textureIsFramebuffer) {
		widthFactor = (float)gstate.getTextureWidth(0) / (float)gstate_c.curTextureWidth;
		heightFactor = (float)gstate.getTextureHeight(0) / (float)gstate_c.curTextureHeight;
	} else {
		widthFactor = 1.0f;
		heightFactor = 1.0f;
	}
	uvScaleOff[0] = widthFactor;
	uvScaleOff[1] = heightFactor;
	uvScaleOff[2] = 0.0f;
	uvScaleOff[3] = 0.0f;
}