package main

import rl "libraries/raylib"
import ui "libraries/snths_ui"
import "core:fmt"
import "core:math/linalg"
import "core:strconv"

UiContainerType :: enum {
	Normal,
	Matrix,
	Scroll,
}

UiElementData_Container :: struct {
	type: UiContainerType,
	thickness: f32,
	scroll: f32,
}

UiElementData_Slider :: struct {
	min, max: f32,
	value:    ^f32,
}

UiElementData_FloatInput :: struct {
	value: ^f32,
	value_string: cstring,
}

UiElementData_Dropdown :: struct {
	value: ^i32,
}

UiButtonType :: enum {
	AddIdentity,
	AddTranslation,
	AddScale,
	AddRotation,
	AddTranslationMatrixConfirm,
	AddScaleMatrixConfirm,
	AddRotationMatrixConfirm,
	ExitDialog,
}

UiElementData_Button :: struct {
	type: UiButtonType,
	options: cstring,
}

UI_ALLOWED_FOCUS :: bit_set[ui.ElementType]{.InputFloat}

MatrixTypes :: enum {
	Identity,
	Translation,
	Scale,
	Rotation,
}

@rodata matrix_add_hotkeys := [MatrixTypes]rl.KeyboardKey {
	.Identity = rl.KeyboardKey.I,
	.Translation = rl.KeyboardKey.T,
	.Scale = rl.KeyboardKey.S,
	.Rotation = rl.KeyboardKey.R,
}

FontStyle :: struct {
	font: rl.Font,
	size, spacing: f32
}

ProgramState :: enum {
	Normal,
	Message,
	Dialog,
}

MouseState :: enum {
	Hover,
	DraggingGrid,
	HoverMatrices,
}

ProgramDialogType :: enum {
	CreateTranslationMatrix,
	CreateScaleMatrix,
	CreateRotationMatrix,
}

@rodata program_dialog_messages := [ProgramDialogType]cstring {
	.CreateTranslationMatrix = "Create tranlation matrix",
	.CreateScaleMatrix = "Create scale matrix",
	.CreateRotationMatrix = "Create rotation matrix",
}

ProgramMessageType :: enum {
	ChangeMatrixTypeClearData,
}

@rodata program_messages := [ProgramMessageType]cstring {
	.ChangeMatrixTypeClearData = "Changing the matrix type will delete all current matrices.\nAre you sure you want to do this?",
}

program: Program

InputFloatField :: [32]byte
MATRIX_POOL_SIZE :: 4096

DIALOG_STATE_INPUT_FIELD_COUNT :: 16
ProgramData_DialogState :: struct {
	input_value_strings: [DIALOG_STATE_INPUT_FIELD_COUNT]InputFloatField,
	input_values: [DIALOG_STATE_INPUT_FIELD_COUNT]f32,
}

Program :: struct {
	ui_context: ui.Context,
	ui_element_focus: ui.ElementId,
	ui_overlay_context: ui.Context,
	ui_overlay_element_focus: ui.ElementId,
	ui_overlay_queue_focus_first: bool,
	matrix_float_pool: []f32,
	matrix_value_string_pool: []InputFloatField,
	matrix_create_dropdown_value: MatrixTypes,
	matrix_count: int,
	matrix_scroll: f32,
	matrix_scroll_speed: f32,
	mouse_state: MouseState,
	grid: Grid,
	font_style: FontStyle,
	state: ProgramState,
	message_type: ProgramMessageType,
	dialog_type: ProgramDialogType,
	dialog_data: ProgramData_DialogState,
}

Grid :: struct {
	area:        rect2,
	inner_area:  rect2,
	zoom:        f32,
}

matrix_pool_max :: proc() -> int {
	return len(program.matrix_float_pool) / 9
}

matrix_pool_get_mat3 :: proc() -> ^mat3 {
	assert(program.matrix_count + 1 < matrix_pool_max(), "Max matrix count exceeded")
	mat := cast(^mat3)&program.matrix_float_pool[program.matrix_count * 9]
	program.matrix_count += 1
	return mat
}

matrix_pool_get_mat3_2 :: proc() -> []InputFloatField {
	assert(program.matrix_count + 1 < matrix_pool_max(), "Max matrix count exceeded")
	value_strings := program.matrix_value_string_pool[program.matrix_count * 9 : program.matrix_count * 9 + 9]
	for &str in value_strings {
		str[0] = '0'
	}
	program.matrix_count += 1
	return value_strings
}

program_show_dialog :: proc(program: ^Program, dialog_type: ProgramDialogType) {
	program.dialog_data = {}
	program.ui_overlay_queue_focus_first = true
	program.state = .Dialog
	program.dialog_type = dialog_type
}

program_show_mesasge :: proc(program: ^Program, message: ProgramMessageType) {
	program.state = .Message
	program.message_type = message
}

program_init :: proc(allocator := context.allocator) {
	program.matrix_float_pool = make(type_of(program.matrix_float_pool), MATRIX_POOL_SIZE)
	program.matrix_value_string_pool = make(type_of(program.matrix_value_string_pool), MATRIX_POOL_SIZE)
	program.matrix_scroll_speed = 32.0
	program.font_style = {
		font = rl.GetFontDefault(),
		size = 16.0,
		spacing = 1.0
	}
	program.grid = {
		area = {0, 0, WINDOW_WIDTH, WINDOW_HEIGHT},
		inner_area = {0, 0, WINDOW_WIDTH, WINDOW_HEIGHT},
		zoom = 1.0,
	}

	ui.init(&program.ui_context, 4000, ui_measure_text, &program.font_style, allocator)
	program.ui_element_focus = ui.element_id_invalid()
	ui.init(&program.ui_overlay_context, 200, ui_measure_text, &program.font_style, allocator)
	program.ui_overlay_element_focus = ui.element_id_invalid()
}

program_update :: proc() {
	ui_elements, matrix_container := program_build_workspace_ui(&program.ui_context)
	ui_overlay_elements: []ui.Element

	switch program.state {
	case .Normal:
		program.mouse_state = program_handle_mouse_workspace(program.mouse_state, matrix_container)
		if rl.IsKeyPressed(.TAB) {
			program.ui_element_focus = ui.focus_next_element(program.ui_element_focus, ui_elements, UI_ALLOWED_FOCUS)
		}
		for hotkey, index in matrix_add_hotkeys {
			if rl.IsKeyPressed(hotkey) {
				switch index {
				case .Identity:    program_add_identity_matrix()
				case .Translation: program_show_dialog(&program, .CreateTranslationMatrix)
				case .Scale:       program_show_dialog(&program, .CreateScaleMatrix)
				case .Rotation:    program_show_dialog(&program, .CreateRotationMatrix)
				}
			}
		}
	case .Message:

	case .Dialog:
		ui_overlay_elements = program_build_dialog_ui(&program.ui_overlay_context, program.dialog_type)
		if rl.IsKeyPressed(.TAB) {
			program.ui_overlay_element_focus = ui.focus_next_element(program.ui_overlay_element_focus, ui_overlay_elements, UI_ALLOWED_FOCUS)
		}
		if rl.IsKeyDown(.ENTER) {
			vals := program.dialog_data.input_values
			switch program.dialog_type {
			case .CreateTranslationMatrix: program_add_translation_matrix(vals[0], vals[1], vals[2])
			case .CreateScaleMatrix:       program_add_scalar_matrix(vals[0], vals[1], vals[2])
			case .CreateRotationMatrix:    program_add_rotation_matrix(vals[0])
			}
			program.state = .Normal
		}
	}

	if program.ui_overlay_queue_focus_first {
		program.ui_overlay_queue_focus_first = false
		program.ui_overlay_element_focus = ui.focus_next_element({container_index = 0, element_index = 0}, ui_overlay_elements, UI_ALLOWED_FOCUS)
	}

	grid := program.grid
	grid_inner_position := rect_get_tl(grid.inner_area)
	draw_grid(grid.area, {32.0, 32.0}, rl.GRAY)

	ui_element_focus := program.state == .Normal ? program.ui_element_focus : ui.element_id_invalid()
	program.ui_element_focus = program_render_ui_and_handle_focus(&program.ui_context, ui_element_focus, program.font_style, ui_elements)

	if len(ui_overlay_elements) > 0 {
		rl.DrawRectangle(0, 0, WINDOW_WIDTH, WINDOW_HEIGHT, {0, 0, 0, 192})
		program.ui_overlay_element_focus = program_render_ui_and_handle_focus(
			&program.ui_overlay_context,
			program.ui_overlay_element_focus,
			program.font_style,
			ui_overlay_elements)
	}
}

program_add_identity_matrix :: proc() {
	value_strings := matrix_pool_get_mat3_2()
	value_strings[0][0] = '1'
	value_strings[4][0] = '1'
	value_strings[8][0] = '1'
}

program_add_translation_matrix :: proc(x, y, z: f32) {
	value_strings := matrix_pool_get_mat3_2()
	fmt.bprint(value_strings[2][:], x)
	fmt.bprint(value_strings[5][:], y)
	fmt.bprint(value_strings[8][:], z)
}

program_add_scalar_matrix :: proc(x, y, z: f32) {
	value_strings := matrix_pool_get_mat3_2()
	fmt.bprint(value_strings[0][:], x)
	fmt.bprint(value_strings[4][:], y)
	fmt.bprint(value_strings[8][:], z)
}

program_add_rotation_matrix :: proc(angle: f32) {
	value_strings := matrix_pool_get_mat3_2()
	rotmat := linalg.matrix2_rotate_f32(angle)
	fmt.bprint(value_strings[0][:], rotmat[0, 0])
	fmt.bprint(value_strings[1][:], rotmat[0, 1])
	fmt.bprint(value_strings[3][:], rotmat[1, 0])
	fmt.bprint(value_strings[4][:], rotmat[1, 1])
}

ui_labeled_input_float_element :: proc(label: cstring, margin: f32, value: ^f32, value_string: ^InputFloatField, allocator := context.temp_allocator) -> ^ui.Element {
	label := fmt.ctprintf("%s: ", label)
	label_element := ui.element(.Label, {}, {label = label, color = ui.col(rl.WHITE), margin = margin})
	dimensions := ui.Dimensions{ui.Size_Pixels(64.0), ui.dimensions_pixels(ui.rect_get_size(ui.get_element_outer_area(label_element^))).height}
	ui.same_line()
	input_data := new(UiElementData_FloatInput, context.temp_allocator)
	input_data^ = {value = value, value_string = cstring(&value_string[0])}
	return ui.element(.InputFloat, dimensions, {margin = margin}, input_data)
}

program_build_dialog_ui :: proc(ctx: ^ui.Context, dialog_type: ProgramDialogType) -> []ui.Element {
	size := v2{WINDOW_WIDTH, WINDOW_HEIGHT}
	container_margin := size * 0.33
	element_margin: f32 = 4.0
	dialog_data := &program.dialog_data

	dialog_type_to_button_type := [ProgramDialogType]UiButtonType {
		.CreateTranslationMatrix = .AddTranslationMatrixConfirm,
		.CreateScaleMatrix = .AddScaleMatrixConfirm,
		.CreateRotationMatrix = .AddRotationMatrixConfirm,
	}
	button_type := dialog_type_to_button_type[program.dialog_type]

	ui.start(ctx, V2_ZERO, {WINDOW_WIDTH, WINDOW_HEIGHT})
	ui.container_begin(ui.dimensions_auto(), {margin = container_margin, padding = 8.0})
		switch dialog_type {
		case .CreateTranslationMatrix, .CreateScaleMatrix:
			ui.element(.Label, {}, {label = program_dialog_messages[program.dialog_type], color = ui.col(rl.WHITE), margin = element_margin})
			ui_labeled_input_float_element("X", 4.0, &dialog_data.input_values[0], &dialog_data.input_value_strings[0])
			ui_labeled_input_float_element("Y", 4.0, &dialog_data.input_values[1], &dialog_data.input_value_strings[1])
			ui_labeled_input_float_element("Z", 4.0, &dialog_data.input_values[2], &dialog_data.input_value_strings[2])

		case .CreateRotationMatrix:
			ui.element(.Label, {}, {label = program_dialog_messages[program.dialog_type], color = ui.col(rl.WHITE), margin = element_margin})
			ui_labeled_input_float_element("Angle", 4.0, &dialog_data.input_values[0], &dialog_data.input_value_strings[0])
		}
		button_data: ^UiElementData_Button
		button_data = new(UiElementData_Button, context.temp_allocator)
		button_data^ = {type = button_type}
		ui.element(.Button, ui.dimensions_pixels({96.0, 32.0}), {label = "Add", margin = element_margin}, button_data)

		button_data = new(UiElementData_Button, context.temp_allocator)
		button_data^ = {type = .ExitDialog}
		ui.same_line()
		ui.element(.Button, ui.dimensions_pixels({96.0, 32.0}), {label = "Cancel", margin = element_margin}, button_data)
	ui.container_end()
	return ui.end()
}

program_build_workspace_ui :: proc(ctx: ^ui.Context) -> (elements: []ui.Element, matrix_container: ^ui.Element) {
	ui.start(ctx, {0, 0}, {WINDOW_WIDTH, WINDOW_HEIGHT})
	top_bar_container_config := ui.ElementConfiguration {margin = 4.0}
	top_bar_element_margin: f32 = 4.0
	top_bar_container := ui.container_begin({ui.Size_Percentage(1.0), ui.Size_Pixels(48.0)}, top_bar_container_config)
	top_bar_inner_height := ui.get_element_inner_area(top_bar_container^).height
	top_bar_button_dimensions := ui.Dimensions {ui.Size_Pixels(top_bar_inner_height), ui.Size_Pixels(top_bar_inner_height)}
	top_bar_button_config := ui.ElementConfiguration {margin = top_bar_element_margin}
	top_bar_dropdown_dimensions := ui.Dimensions {ui.Size_Pixels(160.0), ui.Size_Percentage(1.0)}
	top_bar_dropdown_config := ui.ElementConfiguration {margin = top_bar_element_margin}
	{
		button_config := top_bar_button_config
		button_dimensions := top_bar_button_dimensions
		button_data: ^UiElementData_Button

		button_data = new(UiElementData_Button, context.temp_allocator)
		button_config.label = "I"
		button_data^ = {type = .AddIdentity}
		ui.element(.Button, button_dimensions, button_config, button_data)

		button_data = new(UiElementData_Button, context.temp_allocator)
		button_config.label = "T"
		button_data^ = {type = .AddTranslation}
		ui.same_line()
		ui.element(.Button, button_dimensions, button_config, button_data)

		button_data = new(UiElementData_Button, context.temp_allocator)
		button_config.label = "S"
		button_data^ = {type = .AddScale}
		ui.same_line()
		ui.element(.Button, button_dimensions, button_config, button_data)

		button_data = new(UiElementData_Button, context.temp_allocator)
		button_config.label = "R"
		button_data^ = {type = .AddRotation}
		ui.same_line()
		ui.element(.Button, button_dimensions, button_config, button_data)
	}
	ui.container_end()

	side_bar_container_config := ui.ElementConfiguration {margin = 4.0}
	side_bar_container_data := new(UiElementData_Container, context.temp_allocator)
	side_bar_container_data^ = {type = .Scroll, scroll = program.matrix_scroll}
	matrix_container = ui.container_begin({ui.Size_Percentage(0.2), ui.Size_PercentageRemainder(1.0)}, side_bar_container_config, side_bar_container_data)
	for i in 0..<program.matrix_count {
		container_config := ui.ElementConfiguration{margin = 8.0, label = fmt.ctprintf("Matrix %d", i)}
		container_dimensions := ui.Dimensions {ui.Size_Percentage(1.0), ui.Size_Pixels(112.0)}
		mat := program.matrix_float_pool[i * 9 : i * 9 + 9]
		value_strings := program.matrix_value_string_pool[i * 9 : i * 9 + 9]
		ui_mat3(mat, value_strings, container_dimensions, container_config)
	}
	ui.container_end()
	elements = ui.end()
	return
}

program_handle_mouse_workspace :: proc(mouse_state: MouseState, matrix_container: ^ui.Element) -> MouseState {
	mouse := rl.GetMousePosition()
	switch mouse_state {
	case .Hover:
		if rl.IsMouseButtonPressed(.LEFT) {
			return .DraggingGrid
		} else if rect_has_point(rect2(matrix_container.area), mouse) {
			return .HoverMatrices
		}
	case .DraggingGrid:
		program.grid.inner_area = rect_translate(program.grid.inner_area, rl.GetMouseDelta())
		if !rl.IsMouseButtonDown(.LEFT) {
			return .Hover
		}
	case .HoverMatrices:
		program.matrix_scroll = max(0, program.matrix_scroll - rl.GetMouseWheelMove() * program.matrix_scroll_speed)
		if !rect_has_point(rect2(matrix_container.area), mouse) {
			return .Hover
		}
	}
	return mouse_state
}

program_render_ui_and_handle_focus :: proc(ctx: ^ui.Context, focus_element_id: ui.ElementId, font_style: FontStyle, elements: []ui.Element) -> (focus: ui.ElementId) {
	scissor_container: ^ui.Element
	render_above: ui.Element
	render_above.id = ui.element_id_invalid()
	focus_element_id := focus_element_id

	for &element in elements {
		offset: v2
		if scissor_container != nil {
			if element.id.container_index < scissor_container.id.container_index {
				rl.EndScissorMode()
				scissor_container = nil
			} else {
				offset.y = -(cast(^UiElementData_Container)scissor_container.data).scroll
			}
		}

		area := rect_translate(rl.Rectangle(element.area), offset)
		color := rl.Color(element.color)
		#partial switch element.type {
		case .Container:
			dim := rl.Color {0, 0, 0, 192}
			if element.data != nil {
				data := cast(^UiElementData_Container)element.data
				switch data.type {
				case .Normal:
					rl.DrawRectangleRounded(area, 0.1, 4.0, dim)
					rl.GuiGroupBox(area, element.label)
				case .Matrix:
					matrix_area := ui.rect_grow(ui.rect2(area), element.margin / 2.0)
					tl, tr, br, bl := ui.rect_get_corners(matrix_area)
					col := rl.WHITE
					hook_length: f32 = 4.0
					rl.DrawLineEx(tl, bl, data.thickness, col)
					rl.DrawLineEx(tl, tl + V2_RIGHT * hook_length, data.thickness, col)
					rl.DrawLineEx(bl, bl + V2_RIGHT * hook_length, data.thickness, col)
					rl.DrawLineEx(tr, br, data.thickness, col)
					rl.DrawLineEx(tr, tr + V2_LEFT * hook_length, data.thickness, col)
					rl.DrawLineEx(br, br + V2_LEFT * hook_length, data.thickness, col)
				case .Scroll:
					rl.DrawRectangleRounded(area, 0.1, 4.0, dim)
					rl.GuiGroupBox(area, element.label)
					scissor_container = &element
					rl.BeginScissorMode(**rect2_to_irect2(area))
				}
			} else {
				rl.DrawRectangleRounded(area, 0.1, 4.0, dim)
				rl.GuiGroupBox(area, element.label)
			}
		case .Button:
			if (rl.GuiButton(area, element.label)) {
				if element.data == nil {
					break
				}
				data := cast(^UiElementData_Button)element.data
				switch data.type {
				case .AddIdentity:    program_add_identity_matrix()
				case .AddTranslation: program_show_dialog(&program, .CreateTranslationMatrix)
				case .AddScale:       program_show_dialog(&program, .CreateScaleMatrix)
				case .AddRotation:    program_show_dialog(&program, .CreateRotationMatrix)
				case .AddTranslationMatrixConfirm:
					inputs := program.dialog_data.input_values
					program_add_translation_matrix(inputs[0], inputs[1], inputs[2])
					program.state = .Normal
				case .AddScaleMatrixConfirm:
					inputs := program.dialog_data.input_values
					program_add_scalar_matrix(inputs[0], inputs[1], inputs[2])
					program.state = .Normal
				case .AddRotationMatrixConfirm:
					program_add_rotation_matrix(program.dialog_data.input_values[0])
					program.state = .Normal
				case .ExitDialog:
					program.state = .Normal
				}
			}
		case .Slider:
			data := cast(^UiElementData_Slider)element.data
			rl.GuiSlider(area, element.label, "", data.value, data.min, data.max)
		case .InputFloat:
			data := cast(^UiElementData_FloatInput)element.data
			if rl.GuiValueBoxFloat(area, element.label, data.value_string, data.value, focus_element_id == element.id) == 1 {
				focus_element_id = element.id
			}
		case .Dropdown:
			assert(element.data != nil, "Element data cannot be nil when creating a dropdown")
			data := cast(^UiElementData_Dropdown)element.data
			if focus_element_id == element.id {
				render_above = element
			} else {
				if rl.GuiDropdownBox(area, element.label, data.value, focus_element_id == element.id) {
					if focus_element_id == element.id {
						focus_element_id = ui.element_id_invalid()
					} else {
						focus_element_id = element.id
					}
				}
			}
		case .Label:
			rl.DrawTextEx(font_style.font, element.label, rect_get_tl(area), font_style.size, font_style.spacing, color)
		case:
			rl.DrawRectangleRec(area, rl.MAGENTA)
		}
	}

	if scissor_container != nil {
		rl.EndScissorMode()
		scissor_container = nil
	}

	if render_above.id != ui.element_id_invalid() {
		#partial switch render_above.type {
		case .Dropdown:
			data := cast(^UiElementData_Dropdown)render_above.data
			if rl.GuiDropdownBox(rl.Rectangle(render_above.area), render_above.label, data.value, program.ui_element_focus == render_above.id) {
				if program.ui_element_focus == render_above.id {
					program.ui_element_focus = ui.element_id_invalid()
				} else {
					program.ui_element_focus = render_above.id
				}
			}
		}
	}

	return focus_element_id
}

ui_measure_text :: proc(str: cstring, font_data: rawptr) -> ui.v2 {
	style := cast(^FontStyle)font_data
	return rl.MeasureTextEx(style.font, str, style.size, style.spacing)
}

ui_mat3 :: proc(mat: []f32, value_strings: []InputFloatField, dimensions: ui.Dimensions, config: ui.ElementConfiguration, allocator := context.temp_allocator) {
	container_data := new(UiElementData_Container, allocator)
	container_data^ = {type = .Matrix, thickness = 2.0}
	ui.container_begin(dimensions, config, container_data)
		matrix_input_config := ui.ElementConfiguration {margin = 4.0}
		matrix_input_dimensions := ui.Dimensions {ui.Size_Percentage(1.0 / 3.0), ui.Size_Percentage(1.0 / 3.0)}
		#unroll for i in 0..<3 {
			#unroll for j in 0..<3 {
				data := new(UiElementData_FloatInput, allocator)
				index := i * 3 + j
				data^ = {value = &mat[index], value_string = cstring(&value_strings[index][0])}
				ui.element(.InputFloat, matrix_input_dimensions, matrix_input_config, data)
				ui.same_line()
			}
			program.ui_context.same_line = false
		}
	ui.container_end()
}

draw_grid :: proc(area: rect2, line_interval: v2, color: rl.Color) {
	br := rect_get_br(area)
	for x := area.x; x < br.x; x += line_interval.x {
		rl.DrawLineV({x, area.y}, {x, br.y}, color)
	}
	for y := area.y; y < br.y; y += line_interval.y {
		rl.DrawLineV({area.x, y}, {br.x, y}, color)
	}
}