package game

import "core:log"
import "core:math"
import lg "core:math/linalg"
import "game:assets"
import rl "libs:raylib"

// Some fluid simulation tests

Vec2i32 :: [2]i32

Fluid :: struct {
	size:     Vec2i32,
	friction: f32,
	gravity:  f32,
	terrain:  []f32,
	heights:  []f32,
	flow_x:   []f32,
	flow_y:   []f32,
}

fluid_init :: proc(fluid: ^Fluid, size: Vec2i32, friction: f32, gravity: f32) {
	fluid.size = size
	fluid.friction = friction
	fluid.gravity = gravity
	fluid.terrain = make([]f32, size.x * size.y)
	fluid.heights = make([]f32, size.x * size.y)
	fluid.flow_x = make([]f32, (size.x + 1) * size.y)
	fluid.flow_y = make([]f32, (size.y + 1) * size.x)
}

fluid_destroy :: proc(fluid: ^Fluid) {
	delete(fluid.terrain)
	delete(fluid.heights)
	delete(fluid.flow_x)
	delete(fluid.flow_y)
}

fluid_terrain :: #force_inline proc "contextless" (fluid: ^Fluid, x: i32, y: i32) -> ^f32 {
	idx := fluid.size.x * y + x
	return &fluid.terrain[idx]
}

fluid_height :: #force_inline proc "contextless" (fluid: ^Fluid, x: i32, y: i32) -> ^f32 {
	idx := fluid.size.x * y + x
	return &fluid.heights[idx]
}

fluid_flow_x :: #force_inline proc "contextless" (fluid: ^Fluid, x: i32, y: i32) -> ^f32 {
	idx := (fluid.size.x + 1) * y + x
	return &fluid.flow_x[idx]
}

fluid_flow_y :: #force_inline proc "contextless" (fluid: ^Fluid, x: i32, y: i32) -> ^f32 {
	idx := fluid.size.x * y + x
	return &fluid.flow_y[idx]
}

fluid_simulate :: proc "contextless" (fluid: ^Fluid, dt: f32) {
	dt := dt
	dt = 0.2
	cell_size := f32(1)
	dx := cell_size
	dy := cell_size
	area := f32(1)
	mult := 10 * dt * area / cell_size

	friction := math.pow(1.0 - 0.00, dt)

	// Flow acceleration
	{
		for y in 0 ..< fluid.size.y {
			for x in 1 ..< fluid.size.x {
				fluid_flow_x(fluid, x, y)^ =
					fluid_flow_x(fluid, x, y)^ * friction +
					(fluid_height(fluid, x - 1, y)^ +
							fluid_terrain(fluid, x - 1, y)^ -
							fluid_height(fluid, x, y)^ -
							fluid_terrain(fluid, x, y)^) *
						mult
			}
		}

		for y in 1 ..< fluid.size.y {
			for x in 0 ..< fluid.size.x {
				fluid_flow_y(fluid, x, y)^ =
					fluid_flow_y(fluid, x, y)^ * friction +
					(fluid_height(fluid, x, y - 1)^ +
							fluid_terrain(fluid, x, y - 1)^ -
							fluid_height(fluid, x, y)^ -
							fluid_terrain(fluid, x, y)^) *
						mult
			}
		}
	}

	// Flow Scaling to prevent water height going below 0
	{
		for y in 0 ..< fluid.size.y {
			for x in 0 ..< fluid.size.x {
				total_outflow := f32(0)

				total_outflow -= fluid_flow_x(fluid, x, y)^
				total_outflow -= fluid_flow_y(fluid, x, y)^
				total_outflow += fluid_flow_x(fluid, x + 1, y)^
				total_outflow += fluid_flow_y(fluid, x, y + 1)^

				max_outflow := fluid_height(fluid, x, y)^ * dx * dy / dt

				EPS :: f32(0.0001)
				if total_outflow > 0 {
					scale := min(1, max_outflow / total_outflow)

					if fluid_flow_x(fluid, x, y)^ < 0 {
						fluid_flow_x(fluid, x, y)^ *= scale
					}
					if fluid_flow_y(fluid, x, y)^ < 0 {
						fluid_flow_y(fluid, x, y)^ *= scale
					}
					if fluid_flow_x(fluid, x + 1, y)^ > 0 {
						fluid_flow_x(fluid, x + 1, y)^ *= scale
					}
					if fluid_flow_y(fluid, x, y + 1)^ > 0 {
						fluid_flow_y(fluid, x, y + 1)^ *= scale
					}
				}
			}
		}
	}

	// Height update
	{
		factor := dt / dx / dy
		for y in 0 ..< fluid.size.y {
			for x in 0 ..< fluid.size.x {
				fluid_height(fluid, x, y)^ +=
					(fluid_flow_x(fluid, x, y)^ +
						fluid_flow_y(fluid, x, y)^ -
						fluid_flow_x(fluid, x + 1, y)^ -
						fluid_flow_y(fluid, x, y + 1)^) *
					factor
			}
		}
	}
}

Fluid_Scene :: struct {
	fluid:   Fluid,
	texture: rl.Texture,
}

fluid_scene_init :: proc(scene: ^Fluid_Scene) {
	fluid_init(&scene.fluid, {256, 256}, 0.1, 9.8)
}

fluid_scene_destroy :: proc(scene: ^Fluid_Scene) {
	fluid_destroy(&scene.fluid)
	rl.UnloadTexture(scene.texture)
}

fluid_scene_update :: proc(scene: ^Fluid_Scene, dt: f32) {
	if rl.IsMouseButtonDown(rl.MouseButton.LEFT) || rl.IsMouseButtonDown(rl.MouseButton.RIGHT) {
		sim_size := scene.fluid.size
		mouse_pos := rl.GetMousePosition()
		center := Vec2i32{i32(mouse_pos.x), i32(mouse_pos.y)}

		for y in max(center.y - 100, 0) ..< min(center.y + 100, sim_size.y - 1) {
			for x in max(center.x - 100, 0) ..< max(center.x + 100, sim_size.x - 1) {
				offset_from_center := Vec2i32{x, y} - center
				if lg.length2([2]f32{f32(offset_from_center.x), f32(offset_from_center.y)}) <
				   (100 * 100) {

					if rl.IsMouseButtonDown(.LEFT) {
						fluid_height(&scene.fluid, x, y)^ = 20
					}

					// if rl.IsMouseButtonDown(.RIGHT) {
					// 	fluid_terrain(&scene.fluid, x, y)^ += 10
					// }
				}
			}
		}
	}
	fluid_simulate(&scene.fluid, dt)
}

fluid_scene_draw_3d :: proc(assetman: ^assets.Asset_Manager, scene: ^Fluid_Scene) {
	sim_size := scene.fluid.size
	image: rl.Image
	image.format = rl.PixelFormat.UNCOMPRESSED_R32
	image.width = sim_size.x
	image.height = sim_size.y
	image.mipmaps = 1
	image_data := make([]f32, sim_size.x * sim_size.y, context.temp_allocator)
	image.data = rawptr(&image_data[0])

	for y in 0 ..< sim_size.y {
		for x in 0 ..< sim_size.x {
			image_data[y * sim_size.x + x] = fluid_height(&scene.fluid, x, y)^ * 0.01
		}
	}

	{
		if !rl.IsTextureValid(scene.texture) ||
		   scene.texture.width != sim_size.x ||
		   scene.texture.height != sim_size.y {
			rl.UnloadTexture(scene.texture)

			scene.texture = rl.LoadTextureFromImage(image)
		} else {
			rl.UpdateTexture(scene.texture, image.data)
		}
	}

	if !rl.IsTextureValid(scene.texture) || !rl.IsTextureReady(scene.texture) {
		log.warnf("texture is not valid")
	}

	water_shader := assets.get_shader(
		assetman,
		"assets/shaders/water_vs.glsl",
		"assets/shaders/water_ps.glsl",
		assets.Shader_Location_Set{.LightDir, .LightColor, .Ambient},
	)

	light_dir := lg.normalize(rl.Vector3{1, -1, 0})
	ambient := rl.Vector3{0.1, 0.1, 0.1}
	light_color := rl.Vector3{0.816, 0.855, 0.89}
	rl.SetShaderValue(water_shader.shader, water_shader.locations[.LightDir], &light_dir, .VEC3)
	rl.SetShaderValue(water_shader.shader, water_shader.locations[.Ambient], &ambient, .VEC3)
	rl.SetShaderValue(
		water_shader.shader,
		water_shader.locations[.LightColor],
		&light_color,
		.VEC3,
	)

	model := assets.get_model(assetman, "assets/subdiv_plane.obj")
	model.materials[0].shader = water_shader.shader
	model.materials[0].maps[rl.MaterialMapIndex.ALBEDO].texture = scene.texture
	rl.DrawModel(model, {0, 1, 0}, 1, rl.WHITE)
}

fluid_scene_draw_2d :: proc(assetman: ^assets.Asset_Manager, scene: ^Fluid_Scene) {
	rl.DrawTexture(scene.texture, 0, 0, rl.WHITE)
}