#include "app.hpp"
#include "model.hpp"
#include "ui.hpp"
#include "video.hpp"
extern "C" {
#include "plat.h"
#include "memory.h"
}
#include <string.h>
#include <math.h>
#include <stdio.h>

#define video_arena_size (1024 * 1024 * 16)
#define asset_arena_size (1024 * 1024 * 4)
#define ui_arena_size    (1024 * 16)
#define per_frame_memory_size (1024 * 1024)

static float verts[] = {
	 0.5f,  0.5f, 1.0f, 0.0f, 0.0f, 1.0f, 1.0f,
	-0.5f,  0.5f, 0.0f, 1.0f, 0.0f, 0.0f, 1.0f,
	 0.0f, -0.5f, 0.0f, 0.0f, 1.0f, 0.5f, 0.0f
};

static Buffer_Id upload_verts(Device* dev) {
	Buffer_Id stage = dev->create_buffer(
		"temp",
		sizeof verts,
		Buffer_Flags::copy_src |
		Buffer_Flags::cpu_readwrite
	);
	Buffer_Id vbo = dev->create_buffer(
		"autism triangle vbo",
		sizeof verts,
		Buffer_Flags::vertex_buffer |
		Buffer_Flags::copy_dst
	);
	{
		void* mem;
		mem = dev->map_buffer(stage, 0, sizeof verts);
		memcpy(mem, verts, sizeof verts);
		dev->unmap_buffer(stage);
	}
	{
		Context& ctx = dev->acquire();
		ctx.copy(vbo, stage);
		dev->submit(ctx);
	}
	dev->destroy_bufferi(stage);
	return vbo;
}

Texture_Id make_default_texture(Device* dev) {
	unsigned* mem;
	Texture_Id tex;
	Buffer_Id buf = dev->create_buffer(
		"default texture stage",
		4,
		Buffer_Flags::copy_src |
		Buffer_Flags::cpu_readwrite
	);
	mem = (unsigned*)dev->map_buffer(buf, 0, 4);
	mem[0] = 0xffffffff;
	dev->unmap_buffer(buf);
	tex = dev->create_texture(
		"default PBR texture",
		texture_format_rgba8i,
		Texture_Flags::sampleable | Texture_Flags::copy_dst,
		1,
		1,
		1,
		1,
		1,
		buf
	);
	dev->destroy_bufferi(buf);
	return tex;
}

static Sampler_Id create_clamped_linear(Device* dev) {
	Sampler_State s{};
	s.min = Filter_Mode::linear;
	s.mag = Filter_Mode::linear;
	s.mip = Filter_Mode::linear;
	s.address_u = Address_Mode::repeat;
	s.address_v = Address_Mode::repeat;
	return dev->create_sampler("repeated linear", s);
}

struct Orbit_Cam : public Camera {
	bool first_frame;
	v3f target;
	int px, py;
	int pscroll;
	float yrot = 0.0f;
	static constexpr float sense = 5.0f;
	static constexpr float min_dist = 0.3f;

	void init() {
		pscroll = 0;
		first_frame = true;
		target = v3f(0.0f);
		position = v3f(0.0f, 0.0f, -5.0f);
		Camera::init(
			90.0f,
			v3f::normalised(target - position),
			position
		);
	}

	void update(const App& app) {
		float dx = ((float)(px - app.mx) / (float)app.w) * sense;
		float dy = ((float)(py - app.my) / (float)app.h) * sense;
		int dscroll = pscroll - app.scrolly;
		if (app.mjp(mbtn_middle) || app.mjp(mbtn_right))
			dx = dy = 0.0f;
		if (app.mp(mbtn_middle)) {
			v3f left = v3f::cross(forward, v3f(0.0f, 1.0f, 0.0f));
			v3f up = v3f::cross(forward, left);
			float dist = v3f::mag(position - target);
			position += left * v3f(dx) * dist; 
			position += up * v3f(dy) * dist;
			target = position + forward * dist;
			forward = v3f::normalised(target - position);
		}
		if (app.mp(mbtn_right)) {
			m4f rotation = m4f::rotate(
				m4f::identity(),
				forward.z < 0.0f? (float)dy: (float)-dy,
				v3f(1.0f, 0.0f, 0.0f)
			) * m4f::rotate(
				m4f::identity(),
				(float)dx,
				v3f(0.0f, 1.0f, 0.0f)
			);
			v4f pos = v4f(position - target, 0.0f);
			yrot += dy;
			pos = rotation * pos;
			position = target + v3f(pos.x, pos.y, pos.z);
			forward = v3f::normalised(target - position);
		}
		if (dscroll) {
			v3f zp = position - forward * (float)dscroll;
			float dist = v3f::mag(zp - target);
			if (dist < min_dist || v3f::dot(forward, zp - target) > 0.0f)
				zp = target - forward * min_dist;
			position = zp;
		}
		px = app.mx;
		py = app.my;
		pscroll = app.scrolly;
	}
};

struct Sky {
	Staged_Buffer config;
	Buffer_Id vb;
	Texture_Id texture;
	Shader* shader;
	int vert_binding;
	int picture_binding;
	int cfg_binding;

	struct Cbuffer {
		m4f iview;
		m4f iprojection;
	};

	void init(Device* d, Asset_Arena* assets) {
		Buffer_Id stage;
		void* mem;
		Context& ctx = d->acquire();
		float verts[] = {
			-1.0f, -1.0f, 0.0f, 0.0f,
			-1.0f,  3.0f, 0.0f, 2.0f,
			 3.0f, -1.0f, 2.0f, 0.0f
		};
		Texture* tex = (Texture*)assets->load("sky.tex");
		shader = (Shader*)assets->load("sky.csh");
		texture = tex->id;
		stage = d->create_buffer(
			"sky vb stage",
			sizeof verts,
			Buffer_Flags::cpu_readwrite |
			Buffer_Flags::copy_src
		);
		mem = d->map_buffer(stage, 0, sizeof verts);
		memcpy(mem, verts, sizeof verts);
		d->unmap_buffer(stage);
		vb = d->create_buffer(
			"sky vb",
			sizeof verts,
			Buffer_Flags::copy_dst |
			Buffer_Flags::vertex_buffer
		);
		ctx.copy(vb, stage);
		d->submit(ctx);
		d->destroy_bufferi(stage);
		assert(shader != 0);
		vert_binding = shader->binding_index("verts");
		assert(vert_binding >= 0);
		picture_binding = shader->descriptor_binding("picture");
		assert(picture_binding >= 0);
		cfg_binding = shader->descriptor_binding("config_buffer");
		assert(cfg_binding >= 0);
		config.init(
			d,
			"Sky Cbuffer",
			sizeof(Cbuffer),
			Buffer_Flags::constant_buffer
		);
	}

	void destroy(Device* d) {
		d->destroy_buffer(vb);
		config.destroy(d);
	}

	void render(
		Device* d,
		Arena* a,
		Render_Pass& pass,
		Sampler_Id sampler,
		const Camera& cam
	) {
		Texture_Id backbuffer_id = d->get_backbuffer();
		Texture& backbuffer = d->get_texture(backbuffer_id);
		Cbuffer* cb = (Cbuffer*)config.map(d);
		cb->iview = cam.get_view().inverse();
		cb->iprojection = cam.get_proj().inverse();
		config.unmap(d);
		config.update(d->get_ctx());
		render_imp(
			d,
			a,
			pass,
			sampler,
			config.gpuonly,
			0,
			backbuffer.w,
			backbuffer.h
		);
	}

	void render_imp(
		Device* d,
		Arena* a,
		Render_Pass& pass,
		Sampler_Id sampler,
		Buffer_Id cbuffer,
		int offset,
		int w,
		int h
	) {
		Context& ctx = d->get_ctx();
		Pipeline_Builder pb(a, d);
		pb.begin();
		pb.shader(shader->id);
		pb.vertex_format(shader->vf);
		pb.texture(
			picture_binding,
			texture,
			sampler
		);
		pb.viewport(0, 0, w, h);
		pb.scissor(0, 0, w, h);
		pb.cbuffer(cfg_binding, cbuffer, offset, sizeof(Cbuffer));
		pb.depth(true, false, Depth_Mode::equal);
		Pipeline& pip = pb.build();

		Vertex_Buffer_Binding vbb[] = {{
			.id = vb,
			.offset = 0,
			.target = vert_binding
		}, {}};

		Draw draw{};
		draw.verts = vbb;
		draw.vertex_count = 3;
		draw.instance_count = 1;
		ctx.submit(draw, pip, pass);
	}
};

struct Env_Probe {
	static constexpr int mip_count = 9;
	Texture_Id cubemap[2];
	Texture_Id cubemap_depth;
	Texture_Id faces_render[6];
	Texture_Id faces[6][mip_count];
	Sampler_Id sampler;
	int cubemap_res;
	Staged_Buffer config;
	Staged_Buffer face_cb[6][mip_count];
	Shader* shader;
	int src_binding, vert_binding, face_binding;

	struct Face_Cbuffer {
		float dir[3];
		float roughness;
		float up[3];
	};
	v3f face_dirs[6] = {
		v3f( 1.0f,  0.0f,  0.0f),
		v3f(-1.0f,  0.0f,  0.0f),
		v3f( 0.0f,  1.0f,  0.0f),
		v3f( 0.0f, -1.0f,  0.0f),
		v3f( 0.0f,  0.0f, -1.0f),
		v3f( 0.0f,  0.0f,  1.0f)
	};
	v3f face_ups[6] = {
		v3f(0.0f, 1.0f, 0.0f),
		v3f(0.0f, 1.0f, 0.0f),
		v3f(0.0f, 0.0f, 1.0f),
		v3f(0.0f, 0.0f, -1.0f),
		v3f(0.0f, 1.0f, 0.0f),
		v3f(0.0f, 1.0f, 0.0f)
	};

	void init(Device* dev, Asset_Arena* assets, int res) {
		int i;
		Sampler_State s{};
		cubemap_res = res;
		s.min = Filter_Mode::linear;
		s.mag = Filter_Mode::linear;
		s.mip = Filter_Mode::linear;
		s.address_u = Address_Mode::repeat;
		s.address_v = Address_Mode::repeat;
		s.address_w = Address_Mode::repeat;
		sampler = dev->create_sampler("env probe sampler", s);
		cubemap_depth = dev->create_texture(
			"env cubemap depth",
			texture_format_d16,
			Texture_Flags::sampleable | Texture_Flags::depth_stencil_target,
			res,
			res,
			1,
			1,
			1,
			0
		);
		cubemap[0] = dev->create_texture(
			"env cubemap A",
			texture_format_rgba16f,
			Texture_Flags::cubemap |
			Texture_Flags::colour_target |
			Texture_Flags::sampleable,
			res,
			res,
			1,
			mip_count,
			1,
			Buffer_Id(0)
		);
		cubemap[1] = dev->create_texture(
			"env cubemap B",
			texture_format_rgba16f,
			Texture_Flags::cubemap |
			Texture_Flags::colour_target |
			Texture_Flags::sampleable,
			res,
			res,
			1,
			mip_count,
			1,
			Buffer_Id(0)
		);
		for (i = 0; i < 6; i++) {
			int j;
			int w = res;
			int h = res;
			faces_render[i] = dev->alias_texture(
				cubemap[0],
				"env cubemap face A",
				texture_format_rgba16f,
				Texture_Flags::colour_target |
				Texture_Flags::sampleable,
				w,
				h,
				1,
				1,
				1,
				0,
				i
			);
			for (j = 0; j < mip_count; j++) {
				face_cb[i][j].init(
					dev,
					"env probe mip face cbuffer",
					sizeof(Face_Cbuffer),
					Buffer_Flags::constant_buffer
				);
				faces[i][j] = dev->alias_texture(
					cubemap[1],
					"env cubemap face B",
					texture_format_rgba16f,
					Texture_Flags::colour_target |
					Texture_Flags::sampleable,
					w,
					h,
					1,
					1,
					1,
					j,
					i
				);
				w /= 2;
				h /= 2;
			}
		}
		config.init(
			dev,
			"env probe sky cbuffer",
			sizeof(Sky::Cbuffer) * 6,
			Buffer_Flags::constant_buffer
		);
		shader = (Shader*)assets->load("mip_spec.csh");
		src_binding = shader->descriptor_binding("src");
		face_binding = shader->descriptor_binding("c_face");
		vert_binding = shader->binding_index("verts");
		assert(src_binding >= 0);
		assert(vert_binding >= 0);
	}

	void destroy(Device* dev) {
		int i;
		for (i = 0; i < 6; i++) {
			int j;
			for (j = 0; j < mip_count; j++) {
				dev->destroy_texture(faces[i][j]);
				face_cb[i][j].destroy(dev);
			}
			dev->destroy_texture(faces_render[i]);
		}
		dev->destroy_texture(cubemap_depth);
		dev->destroy_sampler(sampler);
		dev->destroy_texture(cubemap[0]);
		dev->destroy_texture(cubemap[1]);
		config.destroy(dev);
	}

	void update_cbuffer(Device* dev) {
		int i;
		Context& ctx = dev->get_ctx();
		Sky::Cbuffer* bufs = (Sky::Cbuffer*)config.map(dev);
		m4f iproj = m4f::pers(
			90.0f,
			1.0f,
			0.1f,
			1000.0f
		).inverse();
		for (i = 0; i < 6; i++) {
			Sky::Cbuffer& buf = bufs[i];
			m4f view = m4f::lookat(v3f(0.0f), face_dirs[i], face_ups[i]);
			buf.iview = view.inverse();
			buf.iprojection = iproj;
		}
		config.unmap(dev);
		config.update(ctx);
	}

	void render(Device* dev, Arena* a, Sky& sky) {
		int i;
		Pipeline_Builder pb(a, dev);
		Context& ctx = dev->get_ctx();
		update_cbuffer(dev);
		for (i = 0; i < 6; i++) {
			int j;
			for (j = 0; j < mip_count; j++) {
				Face_Cbuffer* cb = (Face_Cbuffer*)face_cb[i][j].map(dev);
				cb->dir[0] = face_dirs[i].x;
				cb->dir[1] = face_dirs[i].y;
				cb->dir[2] = face_dirs[i].z;
				cb->roughness = (float)j / (float)(mip_count - 1);
				cb->up[0] = face_ups[i].x;
				cb->up[1] = face_ups[i].y;
				cb->up[2] = face_ups[i].z;
				face_cb[i][j].unmap(dev);
				face_cb[i][j].update(ctx);
			}
			pb.begin_rp();
			pb.rp_target(faces_render[i], Clear_Mode::discard);
			pb.rp_depth_target(cubemap_depth, 1.0f);
			sky.render_imp(
				dev,
				a,
				pb.build_rp(),
				sampler,
				config.gpuonly,
				i * sizeof(Sky::Cbuffer),
				cubemap_res,
				cubemap_res
			);
		}
		for (i = 0; i < 6; i++) {
			int j;
			for (j = 0; j < mip_count; j++) {
				Texture& t = dev->get_texture(faces[i][j]);
				pb.begin_rp();
				pb.rp_target(t.id, Clear_Mode::discard);
				Render_Pass& pass = pb.build_rp();
				pb.begin();
				pb.viewport(0, 0, t.w, t.h);
				pb.scissor(0, 0, t.w, t.h);
				pb.cbuffer(
					face_binding,
					face_cb[i][j].gpuonly
				);
				pb.texture(
					src_binding,
					cubemap[0],
					sampler
				);
				pb.shader(shader->id);
				pb.vertex_format(shader->vf);
				Pipeline& pip = pb.build();

				Vertex_Buffer_Binding vbb[] = {{
					.id = sky.vb,
					.offset = 0,
					.target = vert_binding
				}, {}};
				Draw draw{};
				draw.verts = vbb;
				draw.vertex_count = 3;
				draw.instance_count = 1;
				ctx.submit(draw, pip, pass);
			}
		}
	}

	Texture_Id get_cubemap() {
		return cubemap[1];
	}
};

struct Config_Buffer {
	float offset[2];
};

struct Config_Buffer2 {
	m4f transform;
	m4f projection;
};

struct C2 : public App {
	Arena video_arena, asset_arena, ui_arena;
	Model_Loader model_loader;
	Material_Loader mat_loader;
	Asset_Arena assets;
	Device* dev;
	Shader* shader, * ui_shader;
	Texture* texture;
	Texture* texture2;
	Texture_Id default_texture;
	Model_Instance* monkey, * monkey2;
	Model_Scene scene;
	Orbit_Cam camera;
	Sky sky;
	Env_Probe eprobe;
	Buffer_Id vbo, cbuf;
	Sampler_Id clamped_linear;
	UI* ui;
	UI::Label* fps_label;
	void* per_frame;
	int frame = 0;
	uint8_t r = 0;
	float rot = 0.0f;
	v3f raxis = v3f(0.0f, 1.0f, 0.0);

	void on_init() override {
		running = 1;
		init_arena(
			&video_arena,
			arena_alloc(arena, video_arena_size),
			video_arena_size
		);
		init_arena(
			&asset_arena,
			arena_alloc(arena, asset_arena_size),
			asset_arena_size
		);
		init_arena(
			&ui_arena,
			arena_alloc(arena, ui_arena_size),
			ui_arena_size
		);
		assets.init(&asset_arena, "pack", 128);
		dev = Device::create(&video_arena, this); 
		default_texture = make_default_texture(dev);
		model_loader.init(dev, &assets);
		mat_loader.init(&assets, default_texture);
		register_asset_loader("MODL", &model_loader);
		register_asset_loader("MTRL", &mat_loader);
		shader = (Shader*)assets.load("triangle.csh");
		ui_shader = (Shader*)assets.load("ui.csh");
		texture = (Texture*)assets.load("22.tex");
		texture2 = (Texture*)assets.load("kita.tex");
		cbuf = dev->create_buffer(
			"config buffer",
			sizeof(Config_Buffer),
			Buffer_Flags::constant_buffer |
			Buffer_Flags::cpu_readwrite
		);
		per_frame = heap_alloc(
			dev->heap,
			per_frame_memory_size
		);
		clamped_linear = create_clamped_linear(dev);
		ui = UI::create(dev, this, &ui_arena, ui_shader->id);
		assert(per_frame != 0);
		vbo = upload_verts(dev);
		ui->layout(w, h);
		fps_label = ui->create_element<UI::Label>(ui->root, "");
		scene.init(arena, 32, clamped_linear);
		monkey = scene.instantiate(
			dev,
			(Model*)assets.load("monkey.mdl")
		);
		monkey2 = scene.instantiate(
			dev,
			(Model*)assets.load("monkey.mdl")
		);
		sky.init(dev, &assets);
		eprobe.init(dev, &assets, 256);
		camera.init();
	}

	void on_update() override {
		Arena frame_arena;
		init_arena(&frame_arena, per_frame, per_frame_memory_size);

		dev->begin_frame();

		if (frame % 10 == 0) {
			char buf[32];
			sprintf(buf, "FPS: %g", 1.0f / dt);
			fps_label->set_text(buf);
		}
		ui->update(&frame_arena);

		{
			void* mem;
			mem = dev->map_buffer(cbuf, 0, sizeof(Config_Buffer));
			Config_Buffer* c = (Config_Buffer*)mem;
			c->offset[0] = sinf((float)frame / 10.0f);
			c->offset[1] = cosf((float)frame / 30.0f);
			dev->unmap_buffer(cbuf);
		}

		Pipeline_Builder pb(&frame_arena, dev);
		pb.begin_rp();
		pb.rp_target(dev->get_backbuffer(), { r, 0x00, 0xff, 0xff });
		Render_Pass& pass = pb.build_rp();

		pb.begin();
		pb.shader(shader->id);
		pb.vertex_format(shader->vf);
		pb.texture(
			shader->descriptor_binding("colour_texture"),
			frame % 2? texture->id: texture2->id,
			clamped_linear
		);
		pb.cbuffer(
			shader->descriptor_binding("config_buffer"),
			cbuf
		);
		Pipeline& pip = pb.build();

		Vertex_Buffer_Binding binding[] = {{
			.id = vbo,
			.offset = 0,
			.target = shader->binding_index("verts")
		}, {}};

		Draw draw{};
		draw.verts = binding;
		draw.vertex_count = 3;
		draw.instance_count = 1;

		Context& ctx = dev->get_ctx();

		ctx.debug_push("autism triangle");
		ctx.submit(draw, pip, pass);
		ctx.debug_pop();

		pb.begin_rp();
		pb.rp_depth_target(dev->get_depth_target(), 1.0f);
		dev->get_ctx().submit(pb.build_rp());

		pb.begin_rp();
		pb.rp_depth_target(dev->get_depth_target(), Clear_Mode::restore);
		Render_Pass& depth_prepass = pb.build_rp();

		pb.begin_rp();
		pb.rp_target(dev->get_backbuffer(), Clear_Mode::restore);
		pb.rp_depth_target(dev->get_depth_target(), Clear_Mode::restore);
		Render_Pass& pass2 = pb.build_rp();

		pb.begin_rp();
		pb.rp_target(dev->get_backbuffer(), Clear_Mode::restore);
		pb.rp_depth_target(dev->get_depth_target(), Clear_Mode::restore);
		Render_Pass& sky_pass = pb.build_rp();

		ctx.debug_push("environment cube");
		eprobe.render(dev, &frame_arena, sky);
		ctx.debug_pop();

		Texture& bb = dev->get_texture(dev->get_backbuffer());
		monkey->transform = m4f::translate(
			m4f::identity(),
			v3f(0.0f, 0.0f, 0.0f)
		) * m4f::rotate(
			m4f::identity(),
			rot,
			raxis
		);
		monkey2->transform = m4f::translate(
			m4f::identity(),
			v3f(2.0f, 0.0f, 2.0f)
		) * m4f::rotate(
			m4f::identity(),
			rot * 2.0f,
			v3f(1.0f, 0.0f, 0.0f)
		);
		camera.asp = (float)bb.w / (float)bb.h;
		camera.update(*this);
		scene.update(camera, dev);
		ctx.debug_push("scene");
		ctx.debug_push("depth prepass");
		scene.render(dev, &frame_arena, depth_prepass, 0, clamped_linear);
		ctx.debug_pop();
		ctx.debug_push("forward");
		scene.render(dev, &frame_arena, pass2, eprobe.get_cubemap(), clamped_linear);
		ctx.debug_pop();
		ctx.debug_pop();

		ctx.debug_push("sky");
		sky.render(
			dev,
			&frame_arena,
			sky_pass,
			clamped_linear,
			camera
		);
		ctx.debug_pop();

		ctx.debug_push("ui");
		ui->render(&frame_arena, dev->get_backbuffer());
		ctx.debug_pop();

		r += 10;
		rot += 1.0f * dt;
		frame++;
		dev->present();
	}

	void on_destroy() override {
		scene.destroy(dev);
		sky.destroy(dev);
		eprobe.destroy(dev);
		ui->destroy();
		assets.destroy();
		dev->destroy_texture(default_texture);
		dev->destroy_sampler(clamped_linear);
		dev->destroy_buffer(vbo);
		dev->destroy_buffer(cbuf);
		dev->destroy();
	}

	void on_resize() override {
		ui->layout(w, h);
		dev->on_resize();
	}
};


extern "C" App* entrypoint(Arena* mem) {
	C2* a = (C2*)arena_alloc(mem, sizeof *a);
	new(a) C2();
	a->init("C2", mem);
	return a;
}