import dlib;
import vulkan;

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

struct GameState
{
	RenderState rds;
	u64         frame;
	u64         frame_idx;
}

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;
	DescSet[2]        desc_set_globals;
	DescSet[2]        desc_set_resources;
	PipelineLayout    pipeline_layout_pbr;

	ImageView         placeholder_tex;
	Buffer            globals_buffer;
}

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

struct Model
{
	Buffer          v_buf;
	Buffer          i_buf;
	Buffer          s_buf;
	ModelState      state;
	ModelRenderInfo info;
	Vertex[]        v;
	u32[]           idx;
}

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;
		};
	}
}

struct Vertex
{
	Vec4 col;
	Vec4 tangent;
	Vec3 pos;
	Vec3 normal;
	Vec2 uv;
}

Model g_box;

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

	return g;
}

void
RunCycle(GameState* g)
{
	g.frame_idx = g.frame % 2;
	
	Reset(&g.rds.frame_arenas[g.frame_idx]);

	Renderer* rd = &g.rds.rd;

	BeginFrame(rd);

	BeginRendering(rd);

	PushConstants(rd, g.rds.pipelines[g_box.info.pid], &g_box.info.pc);

	Bind(rd, g.rds.pipelines[g_box.info.pid], [g.rds.desc_set_globals[g.frame_idx], g.rds.desc_set_resources[g.frame_idx]]);

	BindBuffers(rd, &g_box.i_buf, &g_box.v_buf);

	DrawIndexed(rd, cast(u32)g_box.idx.length, 1, 0);

	FinishRendering(rd);

	SubmitAndPresent(rd);

	g.frame += 1;
}

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

	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[3] resource_bindings = [
		{ binding: 0, descriptorType: DT.Image,   descriptorCount: IMG_MAX, stageFlags: SS.All },
		{ binding: 1, descriptorType: DT.Uniform, descriptorCount: BUF_MAX, stageFlags: SS.All },
		{ binding: 2, descriptorType: DT.Uniform, descriptorCount: UNI_MAX, stageFlags: SS.All },
	];

	Attribute[5] attributes = [
		{ binding: 0, location: 0, format: FMT.RGBA_F32, offset: Vertex.col.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.pipeline_layout_pbr   = CreatePipelineLayout(&rds.rd, [rds.desc_layout_globals, rds.desc_layout_resources], PushConst.sizeof);

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

	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:     LoadAssetData(&rds.frame_arenas[0], "shaders/pbr.vert.spv"),
		frag_shader:       LoadAssetData(&rds.frame_arenas[0], "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);
	bool transfer = Transfer(&rds.rd, &rds.globals_buffer, &rds.globals);

	assert(transfer);

	g_box = MakeBox(rds, 6.0, 6.0, Vec4(0.3, 0.4, 0.8, 1.0));

	CreateBuffer(&rds.rd, &g_box.v_buf, BT.Vertex,  g_box.v.length * Vertex.sizeof, false);
	CreateBuffer(&rds.rd, &g_box.i_buf, BT.Index,   g_box.idx.length * u32.sizeof,  false);
	CreateBuffer(&rds.rd, &g_box.s_buf, BT.Uniform, ModelState.sizeof, false);

	transfer  = Transfer(&rds.rd, &g_box.v_buf, g_box.v);
	transfer &= Transfer(&rds.rd, &g_box.i_buf, g_box.idx);
	transfer &= Transfer(&rds.rd, &g_box.s_buf, &g_box.state);

	assert(transfer);

	Write(&rds.rd, rds.desc_set_globals[0], &rds.globals_buffer, 0, DT.Uniform);
	Write(&rds.rd, rds.desc_set_globals[1], &rds.globals_buffer, 0, DT.Uniform);

	Write(&rds.rd, rds.desc_set_resources[0], &g_box.s_buf, 2, 0, DT.Uniform);
	Write(&rds.rd, rds.desc_set_resources[1], &g_box.s_buf, 2, 0, DT.Uniform);
}

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);
	}
}

Model
MakeBox(RenderState* rds, f32 width, f32 height, Vec4 col)
{
	Model box = {
		v:     Alloc!(Vertex)(&rds.perm_arena, 8),
		idx:   Alloc!(u32)(&rds.perm_arena, 36),
		info:  {
			pid: PID.PBRVVVV,
		},
		state: {
			matrix: Mat4Identity(),
		},
	};

	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),
	];

	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[];

	return box;
}

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