@tool
extends Node

const HT_Painter = preload("./painter.gd")
const HTerrain = preload("../../hterrain.gd")
const HTerrainData = preload("../../hterrain_data.gd")
const HT_Logger = preload("../../util/logger.gd")
const HT_Brush = preload("./brush.gd")

const HT_RaiseShader = preload("./shaders/raise.gdshader")
const HT_SmoothShader = preload("./shaders/smooth.gdshader")
const HT_LevelShader = preload("./shaders/level.gdshader")
const HT_FlattenShader = preload("./shaders/flatten.gdshader")
const HT_ErodeShader = preload("./shaders/erode.gdshader")
const HT_Splat4Shader = preload("./shaders/splat4.gdshader")
const HT_Splat16Shader = preload("./shaders/splat16.gdshader")
const HT_SplatIndexedShader = preload("./shaders/splat_indexed.gdshader")
const HT_ColorShader = preload("./shaders/color.gdshader")
const HT_AlphaShader = preload("./shaders/alpha.gdshader")

const MODE_RAISE = 0
const MODE_LOWER = 1
const MODE_SMOOTH = 2
const MODE_FLATTEN = 3
const MODE_SPLAT = 4
const MODE_COLOR = 5
const MODE_MASK = 6
const MODE_DETAIL = 7
const MODE_LEVEL = 8
const MODE_ERODE = 9
const MODE_COUNT = 10

class HT_ModifiedMap:
	var map_type := 0
	var map_index := 0
	var painter_index := 0

signal flatten_height_changed

var _painters : Array[HT_Painter] = []

var _brush := HT_Brush.new()

var _color := Color(1, 0, 0, 1)
var _mask_flag := false
var _mode := MODE_RAISE
var _flatten_height := 0.0
var _detail_index := 0
var _detail_density := 1.0
var _texture_index := 0
var _slope_limit_low_angle := 0.0
var _slope_limit_high_angle := PI / 2.0

var _modified_maps := []
var _terrain : HTerrain
var _logger = HT_Logger.get_for(self)


func _init():
	for i in 4:
		var p := HT_Painter.new()
		# The name is just for debugging
		p.set_name(str("Painter", i))
		#p.set_brush_size(_brush_size)
		p.texture_region_changed.connect(_on_painter_texture_region_changed.bind(i))
		add_child(p)
		_painters.append(p)


func get_brush() -> HT_Brush:
	return _brush


func get_brush_size() -> int:
	return _brush.get_size()


func set_brush_size(s: int):
	_brush.set_size(s)
#	for p in _painters:
#		p.set_brush_size(_brush_size)


func set_brush_texture(texture: Texture2D):
	_brush.set_shapes([texture])
#	for p in _painters:
#		p.set_brush_texture(texture)


func get_opacity() -> float:
	return _brush.get_opacity()


func set_opacity(opacity: float):
	_brush.set_opacity(opacity)


func set_flatten_height(h: float):
	if h == _flatten_height:
		return
	_flatten_height = h
	flatten_height_changed.emit()


func get_flatten_height() -> float:
	return _flatten_height


func set_color(c: Color):
	_color = c


func get_color() -> Color:
	return _color


func set_mask_flag(m: bool):
	_mask_flag = m


func get_mask_flag() -> bool:
	return _mask_flag


func set_detail_density(d: float):
	_detail_density = clampf(d, 0.0, 1.0)


func get_detail_density() -> float:
	return _detail_density


func set_detail_index(di: int):
	_detail_index = di


func set_texture_index(i: int):
	_texture_index = i


func get_texture_index() -> int:
	return _texture_index


func get_slope_limit_low_angle() -> float:
	return _slope_limit_low_angle


func get_slope_limit_high_angle() -> float:
	return _slope_limit_high_angle


func set_slope_limit_angles(low: float, high: float):
	_slope_limit_low_angle = low
	_slope_limit_high_angle = high


func is_operation_pending() -> bool:
	for p in _painters:
		if p.is_operation_pending():
			return true
	return false


func has_modified_chunks() -> bool:
	for p in _painters:
		if p.has_modified_chunks():
			return true
	return false


func get_undo_chunk_size() -> int:
	return HT_Painter.UNDO_CHUNK_SIZE


func commit() -> Dictionary:
	assert(_terrain.get_data() != null)
	var terrain_data = _terrain.get_data()
	assert(not terrain_data.is_locked())
	
	var changes := []
	var chunk_positions : Array
	
	assert(len(_modified_maps) > 0)
	
	for mm in _modified_maps:
		#print("Flushing painter ", mm.painter_index)
		var painter : HT_Painter = _painters[mm.painter_index]
		var info := painter.commit()
		
		# Note, positions are always the same for each map
		chunk_positions = info.chunk_positions
	
		changes.append({
			"map_type": mm.map_type,
			"map_index": mm.map_index,
			"chunk_initial_datas": info.chunk_initial_datas,
			"chunk_final_datas": info.chunk_final_datas
		})
		
		var cs := get_undo_chunk_size()
		for pos in info.chunk_positions:
			var rect = Rect2(pos * cs, Vector2(cs, cs))
			# This will update vertical bounds and notify normal map baker,
			# since the latter updates out of order for preview
			terrain_data.notify_region_change(rect, mm.map_type, mm.map_index, false, true)
	
#	for i in len(_painters):
#		var p = _painters[i]
#		if p.has_modified_chunks():
#			print("Painter ", i, " has modified chunks")
	
	# `commit()` is supposed to consume these chunks, there should be none left
	assert(not has_modified_chunks())
	
	return {
		"chunk_positions": chunk_positions,
		"maps": changes
	}


func set_mode(mode: int):
	assert(mode >= 0 and mode < MODE_COUNT)
	_mode = mode


func get_mode() -> int:
	return _mode


func set_terrain(terrain: HTerrain):
	if terrain == _terrain:
		return
	_terrain = terrain
	# It's important  to release resources here,
	# otherwise Godot keeps modified terrain maps in memory and "reloads" them like that
	# next time we reopen the scene, even if we didn't save it
	for p in _painters:
		p.set_image(null, null)
		p.clear_brush_shader_params()


# This may be called from an `_input` callback.
# Returns `true` if any change was performed.
func paint_input(position: Vector2, pressure: float) -> bool:
	assert(_terrain.get_data() != null)
	var data := _terrain.get_data()
	assert(not data.is_locked())
	
	if not _brush.configure_paint_input(_painters, position, pressure):
		# Sometimes painting may not happen due to frequency options
		return false

	_modified_maps.clear()
	
	match _mode:
		MODE_RAISE:
			_paint_height(data, position, 1.0)

		MODE_LOWER:
			_paint_height(data, position, -1.0)

		MODE_SMOOTH:
			_paint_smooth(data, position)

		MODE_FLATTEN:
			_paint_flatten(data, position)

		MODE_LEVEL:
			_paint_level(data, position)
			
		MODE_ERODE:
			_paint_erode(data, position)

		MODE_SPLAT:
			# TODO Properly support what happens when painting outside of supported index
			# var supported_slots_count := terrain.get_cached_ground_texture_slot_count()
			# if _texture_index >= supported_slots_count:
			# 	_logger.debug("Painting out of range of supported texture slots: {0}/{1}" \
			# 		.format([_texture_index, supported_slots_count]))
			# 	return
			if _terrain.is_using_indexed_splatmap():
				_paint_splat_indexed(data, position)
			else:
				var splatmap_count := _terrain.get_used_splatmaps_count()
				match splatmap_count:
					1:
						_paint_splat4(data, position)
					4:
						_paint_splat16(data, position)

		MODE_COLOR:
			_paint_color(data, position)

		MODE_MASK:
			_paint_mask(data, position)

		MODE_DETAIL:
			_paint_detail(data, position)
			
		_:
			_logger.error("Unknown mode {0}".format([_mode]))

	assert(len(_modified_maps) > 0)
	return true


func _on_painter_texture_region_changed(rect: Rect2, painter_index: int):
	var data := _terrain.get_data()
	if data == null:
		return
	for mm in _modified_maps:
		if mm.painter_index == painter_index:
			# This will tell auto-baked maps to update (like normals).
			data.notify_region_change(rect, mm.map_type, mm.map_index, false, false)
			break


func _paint_height(data: HTerrainData, position: Vector2, factor: float):
	var image := data.get_image(HTerrainData.CHANNEL_HEIGHT)
	var texture := data.get_texture(HTerrainData.CHANNEL_HEIGHT, 0, true)
	
	var mm := HT_ModifiedMap.new()
	mm.map_type = HTerrainData.CHANNEL_HEIGHT
	mm.map_index = 0
	mm.painter_index = 0
	_modified_maps = [mm]

	# When using sculpting tools, make it dependent on brush size
	var raise_strength := 10.0 + float(_brush.get_size())
	var delta := factor * (2.0 / 60.0) * raise_strength
	
	var p : HT_Painter = _painters[0]
	
	p.set_brush_shader(HT_RaiseShader)
	p.set_brush_shader_param("u_factor", delta)
	p.set_image(image, texture)
	p.paint_input(position)


func _paint_smooth(data: HTerrainData, position: Vector2):
	var image := data.get_image(HTerrainData.CHANNEL_HEIGHT)
	var texture := data.get_texture(HTerrainData.CHANNEL_HEIGHT, 0, true)
	
	var mm := HT_ModifiedMap.new()
	mm.map_type = HTerrainData.CHANNEL_HEIGHT
	mm.map_index = 0
	mm.painter_index = 0
	_modified_maps = [mm]

	var p : HT_Painter = _painters[0]
	
	p.set_brush_shader(HT_SmoothShader)
	p.set_brush_shader_param("u_factor", 1.0)
	p.set_image(image, texture)
	p.paint_input(position)


func _paint_flatten(data: HTerrainData, position: Vector2):
	var image := data.get_image(HTerrainData.CHANNEL_HEIGHT)
	var texture := data.get_texture(HTerrainData.CHANNEL_HEIGHT, 0, true)
	
	var mm := HT_ModifiedMap.new()
	mm.map_type = HTerrainData.CHANNEL_HEIGHT
	mm.map_index = 0
	mm.painter_index = 0
	_modified_maps = [mm]

	var p : HT_Painter = _painters[0]
	
	p.set_brush_shader(HT_FlattenShader)
	p.set_brush_shader_param("u_flatten_value", _flatten_height)
	p.set_image(image, texture)
	p.paint_input(position)


func _paint_level(data: HTerrainData, position: Vector2):
	var image := data.get_image(HTerrainData.CHANNEL_HEIGHT)
	var texture := data.get_texture(HTerrainData.CHANNEL_HEIGHT, 0, true)
	
	var mm := HT_ModifiedMap.new()
	mm.map_type = HTerrainData.CHANNEL_HEIGHT
	mm.map_index = 0
	mm.painter_index = 0
	_modified_maps = [mm]

	var p : HT_Painter = _painters[0]
	
	p.set_brush_shader(HT_LevelShader)
	p.set_brush_shader_param("u_factor", (10.0 / 60.0))
	p.set_image(image, texture)
	p.paint_input(position)


func _paint_erode(data: HTerrainData, position: Vector2):
	var image := data.get_image(HTerrainData.CHANNEL_HEIGHT)
	var texture := data.get_texture(HTerrainData.CHANNEL_HEIGHT, 0, true)
	
	var mm := HT_ModifiedMap.new()
	mm.map_type = HTerrainData.CHANNEL_HEIGHT
	mm.map_index = 0
	mm.painter_index = 0
	_modified_maps = [mm]

	var p : HT_Painter = _painters[0]
	
	p.set_brush_shader(HT_ErodeShader)
	p.set_image(image, texture)
	p.paint_input(position)


func _paint_splat4(data: HTerrainData, position: Vector2):
	var image := data.get_image(HTerrainData.CHANNEL_SPLAT)
	var texture := data.get_texture(HTerrainData.CHANNEL_SPLAT, 0, true)
	var heightmap_texture := data.get_texture(HTerrainData.CHANNEL_HEIGHT, 0)
	
	var mm := HT_ModifiedMap.new()
	mm.map_type = HTerrainData.CHANNEL_SPLAT
	mm.map_index = 0
	mm.painter_index = 0
	_modified_maps = [mm]

	var p : HT_Painter = _painters[0]
	var splat := Color(0.0, 0.0, 0.0, 0.0)
	splat[_texture_index] = 1.0;
	p.set_brush_shader(HT_Splat4Shader)
	p.set_brush_shader_param("u_splat", splat)
	_set_slope_limit_shader_params(p, heightmap_texture)
	p.set_image(image, texture)
	p.paint_input(position)


func _paint_splat_indexed(data: HTerrainData, position: Vector2):
	var map_types := [
		HTerrainData.CHANNEL_SPLAT_INDEX, 
		HTerrainData.CHANNEL_SPLAT_WEIGHT
	]
	_modified_maps = []

	var textures := []
	for mode in 2:
		textures.append(data.get_texture(map_types[mode], 0, true))

	for mode in 2:
		var image := data.get_image(map_types[mode])
		
		var mm := HT_ModifiedMap.new()
		mm.map_type = map_types[mode]
		mm.map_index = 0
		mm.painter_index = mode
		_modified_maps.append(mm)

		var p : HT_Painter = _painters[mode]

		p.set_brush_shader(HT_SplatIndexedShader)
		p.set_brush_shader_param("u_mode", mode)
		p.set_brush_shader_param("u_index_map", textures[0])
		p.set_brush_shader_param("u_weight_map", textures[1])
		p.set_brush_shader_param("u_texture_index", _texture_index)
		p.set_image(image, textures[mode])
		p.paint_input(position)


func _paint_splat16(data: HTerrainData, position: Vector2):
	# Make sure required maps are present
	while data.get_map_count(HTerrainData.CHANNEL_SPLAT) < 4:
		data._edit_add_map(HTerrainData.CHANNEL_SPLAT)

	var splats := []
	for i in 4:
		splats.append(Color(0.0, 0.0, 0.0, 0.0))
	splats[_texture_index / 4][_texture_index % 4] = 1.0

	var textures := []
	for i in 4:
		textures.append(data.get_texture(HTerrainData.CHANNEL_SPLAT, i, true))

	var heightmap_texture := data.get_texture(HTerrainData.CHANNEL_HEIGHT, 0)

	for i in 4:
		var image : Image = data.get_image(HTerrainData.CHANNEL_SPLAT, i)
		var texture : Texture = textures[i]
		
		var mm := HT_ModifiedMap.new()
		mm.map_type = HTerrainData.CHANNEL_SPLAT
		mm.map_index = i
		mm.painter_index = i
		_modified_maps.append(mm)

		var p : HT_Painter = _painters[i]

		var other_splatmaps := []
		for tex in textures:
			if tex != texture:
				other_splatmaps.append(tex)
		
		p.set_brush_shader(HT_Splat16Shader)
		p.set_brush_shader_param("u_splat", splats[i])
		p.set_brush_shader_param("u_other_splatmap_1", other_splatmaps[0])
		p.set_brush_shader_param("u_other_splatmap_2", other_splatmaps[1])
		p.set_brush_shader_param("u_other_splatmap_3", other_splatmaps[2])
		_set_slope_limit_shader_params(p, heightmap_texture)
		p.set_image(image, texture)
		p.paint_input(position)


func _paint_color(data: HTerrainData, position: Vector2):
	var image := data.get_image(HTerrainData.CHANNEL_COLOR)
	var texture := data.get_texture(HTerrainData.CHANNEL_COLOR, 0, true)
	
	var mm := HT_ModifiedMap.new()
	mm.map_type = HTerrainData.CHANNEL_COLOR
	mm.map_index = 0
	mm.painter_index = 0
	_modified_maps = [mm]

	var p : HT_Painter = _painters[0]
	
	# There was a problem with painting colors because of sRGB
	# https://github.com/Zylann/godot_heightmap_plugin/issues/17#issuecomment-734001879

	p.set_brush_shader(HT_ColorShader)
	p.set_brush_shader_param("u_color", _color)
	p.set_brush_shader_param("u_normal_min_y", 0.0)
	p.set_brush_shader_param("u_normal_max_y", 1.0)
	p.set_image(image, texture)
	p.paint_input(position)


func _paint_mask(data: HTerrainData, position: Vector2):
	var image := data.get_image(HTerrainData.CHANNEL_COLOR)
	var texture := data.get_texture(HTerrainData.CHANNEL_COLOR, 0, true)
	
	var mm := HT_ModifiedMap.new()
	mm.map_type = HTerrainData.CHANNEL_COLOR
	mm.map_index = 0
	mm.painter_index = 0
	_modified_maps = [mm]

	var p : HT_Painter = _painters[0]
	
	p.set_brush_shader(HT_AlphaShader)
	p.set_brush_shader_param("u_value", 1.0 if _mask_flag else 0.0)
	p.set_image(image, texture)
	p.paint_input(position)


func _paint_detail(data: HTerrainData, position: Vector2):
	var image := data.get_image(HTerrainData.CHANNEL_DETAIL, _detail_index)
	var texture := data.get_texture(HTerrainData.CHANNEL_DETAIL, _detail_index, true)
	var heightmap_texture = data.get_texture(HTerrainData.CHANNEL_HEIGHT, 0)

	var mm := HT_ModifiedMap.new()
	mm.map_type = HTerrainData.CHANNEL_DETAIL
	mm.map_index = _detail_index
	mm.painter_index = 0
	_modified_maps = [mm]

	var p : HT_Painter = _painters[0]
	var c := Color(_detail_density, _detail_density, _detail_density, 1.0)
	
	# TODO Don't use this shader (why?)
	p.set_brush_shader(HT_ColorShader)
	p.set_brush_shader_param("u_color", c)
	_set_slope_limit_shader_params(p, heightmap_texture)
	p.set_image(image, texture)
	p.paint_input(position)


func _set_slope_limit_shader_params(p: HT_Painter, heightmap_texture: Texture):
	p.set_brush_shader_param("u_normal_min_y", cos(_slope_limit_high_angle))
	p.set_brush_shader_param("u_normal_max_y", cos(_slope_limit_low_angle) + 0.001)
	p.set_brush_shader_param("u_heightmap", heightmap_texture)