import dlib;
import vulkan;

const u32 IMG_MAX      = 100;
const u32 BUF_MAX      = 25;
const u32 UNI_MAX      = 50;
const u32 DESC_SET_MAX = 4;

ImageView[256] TEXTURES;
Buffer[256]    MATERIALS;
Buffer[256]    MODEL_STATES;

struct GameState
{
	RenderState rds;
}

struct RenderState
{
	Renderer                 rd;
	Arena[2]                 frame_arenas;
	Arena                    perm_arena;

	PushConst                pc;
	ShaderGlobals            globals;

	Pipeline[PID.Max]        pipelines;
	DescSetLayout            desc_layout_globals;
	DescSetLayout            desc_layout_resources;
	DescSetLayout            desc_layout_state;
	DescSet[2]               desc_set_globals;
	PipelineLayout           pipeline_layout_pbr;

	ImageView[]              textures;
	Buffer[]                 materials;
	Buffer[]                 model_states;
	u64                      itex, imat, istate;

	ImageView                placeholder_tex;
	Buffer                   globals_buffer;
}

struct ShaderGlobals
{
	Vec4 ambient;
	Vec4 diffuse;
	Vec4 specular;
	f32  shininess = 0.0;
	f32  alpha     = 0.0;
}

struct ModelRenderInfo
{
	PushConst  pc;
	PipelineID pid;

	alias pc this;
}

struct ModelState
{
	Mat4 matrix;
}

enum PBRMod : u32
{
	AlbedoValue     = 0x0001,
	AmbientValue    = 0x0002,
	SpecularValue   = 0x0004,
	AlphaValue      = 0x0008,
	AlbedoTexture   = 0x0010,
	AmbientTexture  = 0x0020,
	SpecularTexture = 0x0040,
	AlphaTexture    = 0x0080,
}

enum PipelineID : u32
{
	None,
	PBRVVVV,
	PBRTVVV,
	PBRVTVV,
	PBRVVTV,
	PBRVVVT,
	PBRTTVV,
	PBRTVTV,
	PBRTVVT,
	PBRVTTV,
	PBRVTVT,
	PBRVVTT,
	PBRVTTT,
	PBRTVTT,
	PBRTTVT,
	PBRTTTV,
	PBRTTTT,
	Max,
}

alias PID = PipelineID;

const PID[] PBR_PIPELINES = [
	PID.PBRVVVV, PID.PBRTVVV, PID.PBRVTVV, PID.PBRVVTV, 
	PID.PBRVVVT, PID.PBRTTVV, PID.PBRTVTV, PID.PBRTVVT, 
	PID.PBRVTTV, PID.PBRVTVT, PID.PBRVVTT, PID.PBRTTTV, 
	PID.PBRTTVT, PID.PBRTVTT, PID.PBRVTTT, PID.PBRTTTT,
];

struct PushConst
{
	union
	{
		struct
		{
			u32 t0;
			u32 t1;
			u32 t2;
			u32 t3;
			u32 m0;
			u32 s0;
		};
		struct
		{
			u32 albedo;
			u32 ambient;
			u32 specular;
			u32 alpha;
			u32 material;
			u32 state;
		};
	}
}

ModelData g_box;

void
RunCycle(GameState* g)
{

}

GameState
InitGame(PlatformWindow* window)
{
	GameState g;
	Init(&g.rds, window);

	return g;
}

void
Init(RenderState* rds, PlatformWindow* window)
{
	version(linux)
	{
		PlatformHandles handles = {
			conn:   window.conn,
			window: window.window,
		};
	}

	const u64 resource_max = 256;

	rds.textures     = TEXTURES;
	rds.materials    = MATERIALS;
	rds.model_states = MODEL_STATES;

	DescLayoutBinding[3] global_bindings = [
		{ binding: 0, descriptorType: DT.Uniform,         descriptorCount: 1, stageFlags: SS.All },
		{ binding: 1, descriptorType: DT.StorageTexelBuf, descriptorCount: 1, stageFlags: SS.All },
		{ binding: 2, descriptorType: DT.StorageImage,    descriptorCount: 1, stageFlags: SS.All },
	];

	DescLayoutBinding[6] resource_bindings = [
		{ binding: 0, descriptorType: DT.Image,   descriptorCount: 1, stageFlags: SS.All },
		{ binding: 1, descriptorType: DT.Image,   descriptorCount: 1, stageFlags: SS.All },
		{ binding: 2, descriptorType: DT.Image,   descriptorCount: 1, stageFlags: SS.All },
		{ binding: 3, descriptorType: DT.Image,   descriptorCount: 1, stageFlags: SS.All },
		{ binding: 4, descriptorType: DT.Uniform, descriptorCount: 1, stageFlags: SS.All },
		{ binding: 5, descriptorType: DT.Uniform, descriptorCount: 1, stageFlags: SS.All },
	];

	DescLayoutBinding[1] state_bindings = [
		{ binding: 0, descriptorType: DT.Uniform, descriptorCount: 1, stageFlags: SS.All },
	];

	Attribute[5] attributes = [
		{ binding: 0, location: 0, format: FMT.RGBA_F32, offset: Vertex.color.offsetof   },
		{ binding: 0, location: 1, format: FMT.RGBA_F32, offset: Vertex.tangent.offsetof },
		{ binding: 0, location: 2, format: FMT.RGB_F32,  offset: Vertex.pos.offsetof     },
		{ binding: 0, location: 3, format: FMT.RGB_F32,  offset: Vertex.normal.offsetof  },
		{ binding: 0, location: 4, format: FMT.RG_F32,   offset: Vertex.uv.offsetof      },
	];

	rds.rd           = InitRenderer(handles, MB(24), MB(32));
	rds.perm_arena   = CreateArena(MB(4));
	rds.frame_arenas = [
		CreateArena(MB(4)),
		CreateArena(MB(4)),
	];

	rds.desc_layout_globals   = CreateDescSetLayout(&rds.rd, global_bindings);
	rds.desc_layout_resources = CreateDescSetLayout(&rds.rd, resource_bindings);
	rds.desc_layout_state     = CreateDescSetLayout(&rds.rd, state_bindings);

	rds.pipeline_layout_pbr   = CreatePipelineLayout(&rds.rd, [rds.desc_layout_globals, rds.desc_layout_resources, rds.desc_layout_state], PushConst.sizeof);

	foreach(i; 0 .. 2)
	{
		rds.desc_set_globals[i]   = AllocDescSet(&rds.rd, rds.desc_layout_globals);
	}

	u32[4] spec_data;

	Specialization spec = {
		data: spec_data.ptr,
		size: u32.sizeof * spec_data.length,
		entries: [
			{ constantID: 0, size: u32.sizeof, offset: u32.sizeof*0 },
			{ constantID: 1, size: u32.sizeof, offset: u32.sizeof*1 },
			{ constantID: 2, size: u32.sizeof, offset: u32.sizeof*2 },
			{ constantID: 3, size: u32.sizeof, offset: u32.sizeof*3 },
		],
	};

	GfxPipelineInfo pbr_info = {
		vertex_shader:     LoadFile(&rds.frame_arenas[0], "assets/shaders/pbr.vert.spv"),
		frag_shader:       LoadFile(&rds.frame_arenas[0], "assets/shaders/pbr.frag.spv"),
		input_rate:        IR.Vertex,
		input_rate_stride: Vertex.sizeof,
		layout:            rds.pipeline_layout_pbr,
		vertex_attributes: attributes,
		vert_spec:         spec,
		frag_spec:         spec,
	};

	foreach(pid; PBR_PIPELINES)
	{
		u32 mod      = PIDToPBR(pid);
		spec_data[0] = mod & PBRMod.AlbedoTexture;
		spec_data[1] = mod & PBRMod.AmbientTexture;
		spec_data[2] = mod & PBRMod.SpecularTexture;
		spec_data[3] = mod & PBRMod.AlphaTexture;
		
		bool result = CreateGraphicsPipeline(&rds.rd, &rds.pipelines[pid], &pbr_info);
		assert(result);
	}

	const u64 tex_size = 32*32*4;
	u8[tex_size] placeholder_tex;

	u8[4] magenta = [255, 0, 255, 255];
	u8[4] black   = [0,   0, 0,   255];
	u64 half      = tex_size/2;
	for(u64 i = 0; i < tex_size; i += 32)
	{
		bool swap = i <= half;
		for(u64 j = 0; j < 16; j += 4)
		{
			placeholder_tex[i+j .. i+j+4]       = !swap ? magenta[0 .. $] : black[0 .. $];
			placeholder_tex[i+j+16 .. i+j+16+4] = !swap ? black[0 .. $]   : magenta[0 .. $];
		}
	}

	CreateImageView(&rds.rd, &rds.placeholder_tex, 32, 32, 4, placeholder_tex);

	CreateBuffer(&rds.rd, &rds.globals_buffer, BT.Uniform, ShaderGlobals.sizeof, false);

	ModelData md = LoadGLTF(&rds.frame_arenas[0], "assets/models/DamagedHelmet.glb");
}

PipelineID
PBRToPID(u32 mod)
{
	switch(mod)
	{
		case GetPBRMod(false, false, false, false): return PID.PBRVVVV;
		case GetPBRMod(true , false, false, false): return PID.PBRTVVV;
		case GetPBRMod(false, true , false, false): return PID.PBRVTVV;
		case GetPBRMod(false, false, true , false): return PID.PBRVVTV;
		case GetPBRMod(false, false, false, true ): return PID.PBRVVVT;
		case GetPBRMod(true , true , false, false): return PID.PBRTTVV;
		case GetPBRMod(true , false, true , false): return PID.PBRTVTV;
		case GetPBRMod(true , false, false, true ): return PID.PBRTVVT;
		case GetPBRMod(false, true , true , false): return PID.PBRVTTV;
		case GetPBRMod(false, true , false, true ): return PID.PBRVTVT;
		case GetPBRMod(false, false, true , true ): return PID.PBRVVTT;
		case GetPBRMod(false, true , true , true ): return PID.PBRVTTT;
		case GetPBRMod(true , false, true , true ): return PID.PBRTVTT;
		case GetPBRMod(true , true , false, true ): return PID.PBRTTVT;
		case GetPBRMod(true , true , true , false): return PID.PBRTTTV;
		case GetPBRMod(true , true , true , true ): return PID.PBRTTTT;
		default:                                    return PID.None;
	}
}

u32
PIDToPBR(PipelineID pid)
{
	switch(pid) with(PID)
	{
		case PBRVVVV: return GetPBRMod(false, false, false, false);
		case PBRTVVV: return GetPBRMod(true , false, false, false);
		case PBRVTVV: return GetPBRMod(false, true , false, false);
		case PBRVVTV: return GetPBRMod(false, false, true , false);
		case PBRVVVT: return GetPBRMod(false, false, false, true );
		case PBRTTVV: return GetPBRMod(true , true , false, false);
		case PBRTVTV: return GetPBRMod(true , false, true , false);
		case PBRTVVT: return GetPBRMod(true , false, false, true );
		case PBRVTTV: return GetPBRMod(false, true , true , false);
		case PBRVTVT: return GetPBRMod(false, true , false, true );
		case PBRVVTT: return GetPBRMod(false, false, true , true );
		case PBRTTTV: return GetPBRMod(true , true , true , false);
		case PBRTTVT: return GetPBRMod(true , true , false, true );
		case PBRTVTT: return GetPBRMod(true , false, true , true );
		case PBRVTTT: return GetPBRMod(false, true , true , true );
		case PBRTTTT: return GetPBRMod(true , true , true , true );
		default:      return 0;
	}
}

static u32
GetPBRMod(bool albedo = false, bool ambient = false, bool specular = false, bool alpha = false)
{
	with(PBRMod)
	{
		return 
			(albedo   ? AlbedoTexture   : AlbedoValue)   |
			(ambient  ? AmbientTexture  : AmbientValue)  |
			(specular ? SpecularTexture : SpecularValue) |
			(alpha    ? AlphaTexture    : AlphaValue);
	}
}

/*
ModelData
MakeBox(RenderState* rds, f32 width, f32 height, Vec4 col)
{
	ModelData box = {
		v:     Alloc!(Vertex)(&rds.frame_arenas[0], 8),
		idx:   Alloc!(u32)(&rds.frame_arenas[0], 36),
		parts: Alloc!(MeshPart)(&rds.frame_arenas[0], 1),
		state: {
			matrix: Mat4Identity(),
		},
	};

	static const Vec3[8] positions = [
		Vec3(-1.0, 0.0, -1.0), Vec3(+1.0, 0.0, -1.0), Vec3(+1.0, 1.0, -1.0), Vec3(-1.0, 1.0, -1.0),
		Vec3(-1.0, 0.0, +1.0), Vec3(+1.0, 0.0, +1.0), Vec3(+1.0, 1.0, +1.0), Vec3(-1.0, 1.0, +1.0),
	];

	static const u32[36] indices = [
		0, 1, 3, 3, 1, 2,
    1, 5, 2, 2, 5, 6,
    5, 4, 6, 6, 4, 7,
    4, 0, 7, 7, 0, 3,
    3, 2, 7, 7, 2, 6,
    4, 5, 0, 0, 5, 1,
	];

	f32 half_width = width/2.0;

	Vec3 pos = Vec3(half_width, height, half_width);
	for(u64 i = 0; i < positions.length; i += 1)
	{
		box.v[i].pos = pos * positions[i];
	}

	box.idx[] = indices[];
	
	box.parts[0].mat    = 0;
	box.parts[0].offset = 0;
	box.parts[0].length = indices.length;

	return box;
}

Model
Upload(RenderState* rds, ModelData* data)
{
	Model model;
	u32[] tex_idx   = Alloc!(u32)(&rds.frame_arenas[0], data.text.length);
	u32[] mat_idx   = Alloc!(u32)(&rds.frame_arenas[0], data.materials.length);
	u32[] state_idx = Alloc!(u32)(&rds.frame_arenas[0], data.model_states.length);

	bool result = true;

	result &= Transfer(&rds.rd, &model.v_buf, data.v);
	result &= Transfer(&rds.rd, &model.i_buf, data.idx);

	assert(result);

	for(u64 i = 0; i < data.tex.length; i += 1)
	{
		TextureData* tex = &data.tex[i];
		ImageView* view  = &rds.textures[rds.itex++];
		CreateImageView(&rds.rd, view, tex.width, tex.height, tex.ch, tex.data);
	}

	for(u64 i = 0; i < data.materials.length; i += 1)
	{
		Buffer* buf = &rds.materials[rds.imat++];
		CreateBuffer(&rds.rd, buf, BT.Uniform, Material.sizeof);
		//result = Transfer(&rds.rd, buf, )
	}

	for(u64 i = 0; i < data.model_states.length; i += 1)
	{
		Buffer* buf = &rds.model_states[rds.istate++];
		//CreateBuffer(&rds)
	}

	model.parts = data.parts;
}
*/

unittest
{
	with(PBRMod) with(PID) 
	{
		foreach(pid; PBR_PIPELINES)
		{
			assert(PBRToPID(PIDToPBR(pid)) == pid);
		}
	}
}