millimeters-of-aluminum/addons/module_builder_plugin/module_builder_editor_plugin.gd

@tool
class_name BuilderEditor
extends EditorPlugin

# --- Constants and Scene References ---
const MAIN_EDITOR_SCENE = preload("res://addons/module_builder_plugin/module_editor.tscn")
const BUILDER_DOCK_SCENE = preload("res://addons/module_builder_plugin/builder_dock.tscn")
const CONSTRUCTION_TREE_SCENE = preload("res://addons/module_builder_plugin/construction_tree.tscn")
const CONSTRUCTION_INSPECTOR_SCENE = preload("res://addons/module_builder_plugin/construction_inspector.tscn")

const MODULE_SCENE = preload("res://scenes/ship/builder/module.tscn")

# --- Dock references ---
var builder_dock: Control
var construction_tree_dock: Control
var construction_inspector_dock: Control
var tree_control: Tree

# --- Node References from the main screen scene ---
var builder_world: World2D
var main_screen: Control
var main_viewport: SubViewport
var zoom_label: Label
var rotate_button: Button
var center_button: Button
var pressurize_button: Button
var save_button: Button
var builder_scene_root: Node2D
var builder_camera: Camera2D

# --- State Management Enum ---
enum BuilderState { 
	IDLE, 
	PLACING_STRUCTURAL, 
	PLACING_COMPONENT,
	WIRING # For future use
}
var current_state: BuilderState = BuilderState.IDLE

# --- State Variables ---
var preview_node = null # Can be either StructuralPiece or Component
var active_scene: PackedScene = null # Can be either StructuralPiece or Component
var rotation_angle: float = 0.0
var grid_size: float = 50.0

var undo_redo: EditorUndoRedoManager

# --- Most of the setup functions remain the same ---

func _enter_tree():
	main_screen = MAIN_EDITOR_SCENE.instantiate()
	EditorInterface.get_editor_main_screen().add_child(main_screen)
	
	main_viewport = main_screen.find_child("SubViewport")
	builder_camera = main_screen.find_child("Camera2D")
	
	# Get button and label references
	zoom_label = main_screen.find_child("ZoomLabel")
	rotate_button = main_screen.find_child("RotateButton")
	center_button = main_screen.find_child("CenterButton")
	pressurize_button = main_screen.find_child("PressuriseButton")
	save_button = main_screen.find_child("SaveButton")
	
	_setup_builder_world()
	_setup_button_connections()
	_update_ui_labels()
	
	# Add the Tool Menu Item
	add_tool_menu_item("Generate Structure Definitions", _on_generate_structures_pressed)
	
	main_screen.hide()
	
	undo_redo = EditorInterface.get_editor_undo_redo()
	undo_redo.action_is_committing.connect(_on_undo_redo_action_committed)

func _setup_builder_world():
	builder_world = World2D.new()
	if is_instance_valid(main_viewport):
		main_viewport.world_2d = builder_world
		
		builder_scene_root = Node2D.new()
		builder_scene_root.name = "BuilderRoot"
		main_viewport.add_child(builder_scene_root)

func _setup_docks():
	if BUILDER_DOCK_SCENE:
		builder_dock = BUILDER_DOCK_SCENE.instantiate()
		builder_dock.active_piece_set.connect(on_active_piece_set)
		add_control_to_bottom_panel(builder_dock, "Ship Builder")
		
	if CONSTRUCTION_TREE_SCENE:
		construction_tree_dock = CONSTRUCTION_TREE_SCENE.instantiate()
		tree_control = construction_tree_dock.find_child("Tree")
		add_control_to_dock(DOCK_SLOT_LEFT_UR, construction_tree_dock)
		_refresh_tree_display()
		builder_world.changed.connect(_refresh_tree_display)
	
	
	if CONSTRUCTION_INSPECTOR_SCENE:
		construction_inspector_dock = CONSTRUCTION_INSPECTOR_SCENE.instantiate()
		add_control_to_dock(DOCK_SLOT_RIGHT_UL, construction_inspector_dock)

func switch_to_dock_tab(dock_control: Control, tab_name: String):
	var tab_container = dock_control.find_child("TabContainer")
	if not is_instance_valid(tab_container):
		print("Error: TabContainer not found in dock control.")
		return

	for i in range(tab_container.get_tab_count()):
		if tab_container.get_tab_title(i) == tab_name:
			tab_container.current_tab = i
			return

	print("Warning: Tab '%s' not found." % tab_name)

func _teardown_docks():
	if builder_dock:
		remove_control_from_bottom_panel(builder_dock)
		builder_dock.queue_free()

	if construction_tree_dock:
		remove_control_from_docks(construction_tree_dock)
		construction_tree_dock.queue_free()
		builder_world.changed.disconnect(_refresh_tree_display)
	
	if construction_inspector_dock:
		remove_control_from_docks(construction_inspector_dock)
		construction_inspector_dock.queue_free()

func _exit_tree():
	if main_screen:
		main_screen.queue_free()

	# Clean up the menu item
	remove_tool_menu_item("Generate Structure Definitions")

func _has_main_screen() -> bool:
	return true

func _make_visible(visible):
	if main_screen:
		main_screen.visible = visible
		_setup_gui_input_listener(visible)
		
		if visible:
			_setup_docks()
		else:
			_teardown_docks()

func _get_plugin_name():
	return "Ship Builder"
	
func _get_plugin_icon():
	return EditorInterface.get_editor_theme().get_icon("Node", "EditorIcons")

func _setup_gui_input_listener(connect: bool):
	if main_screen:
		if connect:
			main_screen.gui_input.connect(_on_viewport_input)
		else:
			main_screen.gui_input.disconnect(_on_viewport_input)
			
func _setup_button_connections():
	if rotate_button: rotate_button.pressed.connect(_on_rotate_button_pressed)
	if center_button: center_button.pressed.connect(_on_center_button_pressed)
	if pressurize_button: pressurize_button.pressed.connect(_on_pressurise_button_pressed)
	if save_button: save_button.pressed.connect(_on_save_button_pressed)

func _update_ui_labels():
	if is_instance_valid(zoom_label):
		var zoom_percent = int(builder_camera.zoom.x * 100)
		zoom_label.text = "Zoom: %d%%" % zoom_percent
	
func _process(_delta):
	_update_ui_labels()
	_refresh_tree_display()
	
func _on_viewport_input(event: InputEvent) -> void:
	if event is InputEventMouseMotion and event.button_mask & MOUSE_BUTTON_MASK_RIGHT:
		builder_camera.position -= event.relative / builder_camera.zoom
		
	if event is InputEventMouseButton and event.button_index == MOUSE_BUTTON_WHEEL_UP:
		builder_camera.zoom *= 1.1
	elif event is InputEventMouseButton and event.button_index == MOUSE_BUTTON_WHEEL_DOWN:
		builder_camera.zoom /= 1.1

	if event is InputEventMouseMotion:
		_update_preview_position()

	match current_state:
		BuilderState.IDLE:
			if event is InputEventMouseButton and event.button_index == MOUSE_BUTTON_RIGHT and event.is_pressed():
				_remove_piece_under_mouse()

		BuilderState.PLACING_STRUCTURAL:
			if event is InputEventMouseButton and event.button_index == MOUSE_BUTTON_LEFT and event.is_pressed():
				_place_piece_from_preview()
			if event is InputEventMouseButton and event.button_index == MOUSE_BUTTON_RIGHT and event.is_pressed():
				on_clear_preview()

		BuilderState.PLACING_COMPONENT:
			if event is InputEventMouseButton and event.button_index == MOUSE_BUTTON_LEFT and event.is_pressed():
				_place_component_from_preview()
			if event is InputEventMouseButton and event.button_index == MOUSE_BUTTON_RIGHT and event.is_pressed():
				on_clear_preview()

		BuilderState.WIRING:
			pass


func _unhandled_key_input(event: InputEvent):
	if not event.is_pressed(): return
	if event is InputEventKey and event.as_text().to_lower() == "r":
		_on_rotate_button_pressed()
		get_tree().set_input_as_handled()


func on_active_piece_set(scene: PackedScene):
	if is_instance_valid(preview_node):
		preview_node.queue_free()
		
	current_state = BuilderState.PLACING_STRUCTURAL
	active_scene = scene
	preview_node = scene.instantiate() as StructuralPiece
	preview_node.is_preview = true
	builder_scene_root.add_child(preview_node)
	_update_preview_position()

func _on_component_selected(component_scene: PackedScene):
	if is_instance_valid(preview_node):
		preview_node.queue_free()
	
	current_state = BuilderState.PLACING_COMPONENT
	active_scene = component_scene
	preview_node = component_scene.instantiate() as Component
	builder_scene_root.add_child(preview_node)
	
	print("Now placing component: ", component_scene.resource_path)

func on_clear_preview():
	if is_instance_valid(preview_node):
		preview_node.queue_free()
		preview_node = null
		active_scene = null
	current_state = BuilderState.IDLE


func _update_preview_position():
	if not is_instance_valid(preview_node):
		return
		
	var viewport: SubViewport = main_screen.find_child("SubViewport")
	if not viewport: return
	
	var world_mouse_pos = viewport.get_canvas_transform().affine_inverse() * viewport.get_mouse_position()
	
	match current_state:
		BuilderState.PLACING_STRUCTURAL:
			var snapped_pos = world_mouse_pos.snapped(Vector2(grid_size, grid_size))
			preview_node.global_position = snapped_pos
			preview_node.rotation = rotation_angle
		BuilderState.PLACING_COMPONENT:
			var target_module = _find_first_module()
			if target_module:
				var closest_point = _find_closest_attachment_point(target_module, world_mouse_pos)
				if closest_point:
					preview_node.global_position = closest_point.position
				else:
					preview_node.global_position = world_mouse_pos
			else:
				preview_node.global_position = world_mouse_pos


# --- REFACTORED: Piece Placement ---
func _place_piece_from_preview():
	if not is_instance_valid(preview_node) or not is_instance_valid(active_scene):
		return

	var viewport: SubViewport = main_screen.find_child("SubViewport")
	if not viewport: return
	
	var world_mouse_pos = viewport.get_canvas_transform().affine_inverse() * viewport.get_mouse_position()
	var snapped_pos = world_mouse_pos.snapped(Vector2(grid_size, grid_size))

	var target_module = _find_nearby_modules(snapped_pos)
	if not target_module:
		target_module = MODULE_SCENE.instantiate() as Module
		builder_scene_root.add_child(target_module)
		target_module.global_position = snapped_pos
		target_module.owner = builder_scene_root
	
	var piece_to_place = active_scene.instantiate()
	
	# --- The main change: Add as a direct child of the module ---
	target_module.add_child(piece_to_place)
	piece_to_place.owner = target_module
	piece_to_place.rotation = rotation_angle
	piece_to_place.global_position = snapped_pos
	
	undo_redo.create_action("Place Structural Piece")
	undo_redo.add_do_method(target_module, "add_child", piece_to_place)
	undo_redo.add_do_method(piece_to_place, "set_owner", target_module)
	undo_redo.add_do_method(target_module, "_recalculate_collision_shape")
	undo_redo.add_undo_method(target_module, "remove_child", piece_to_place)
	undo_redo.add_undo_method(target_module, "_recalculate_collision_shape")
	undo_redo.commit_action()

# --- Component Placement remains the same ---
func _place_component_from_preview():
	if not is_instance_valid(preview_node) or not is_instance_valid(active_scene):
		push_error("Cannot place component: Invalid preview or scene.")
		return

	var viewport: SubViewport = main_screen.find_child("SubViewport")
	if not viewport: return

	var world_mouse_pos = viewport.get_canvas_transform().affine_inverse() * viewport.get_mouse_position()
	
	var target_module = _find_first_module()
	if not target_module: 
		push_error("No module found to attach component to.")
		return

	var closest_point = _find_closest_attachment_point(target_module, world_mouse_pos)
	if not closest_point: 
		print("No valid attachment point nearby.")
		return

	var component_to_place = active_scene.instantiate() as Component
	
	target_module.attach_component(component_to_place, closest_point.position, closest_point.piece)
	
	undo_redo.create_action("Place Component")
	undo_redo.add_do_method(target_module, "attach_component", component_to_place, closest_point.position, closest_point.piece)
	undo_redo.add_undo_method(target_module, "remove_child", component_to_place)
	undo_redo.add_do_method(target_module, "recalculate_physical_properties")
	undo_redo.add_undo_method(target_module, "recalculate_physical_properties")
	undo_redo.commit_action()
	
	preview_node.global_position = closest_point.position

# --- Find Nearby Modules remains the same ---
func _find_nearby_modules(position: Vector2) -> Module:
	const OVERLAP_MARGIN = 20.0
	
	if not active_scene or not active_scene.can_instantiate(): return null
	var piece_instance = active_scene.instantiate()
	var shape_node = piece_instance.find_child("CollisionShape2D")
	if not shape_node: 
		piece_instance.queue_free()
		return null
	var piece_shape = shape_node.shape
	piece_instance.queue_free()
	
	var enlarged_shape
	if piece_shape is RectangleShape2D:
		enlarged_shape = RectangleShape2D.new()
		enlarged_shape.size = piece_shape.size + Vector2(OVERLAP_MARGIN, OVERLAP_MARGIN) * 2
	elif piece_shape is CapsuleShape2D:
		enlarged_shape = CapsuleShape2D.new()
		enlarged_shape.radius = piece_shape.radius + OVERLAP_MARGIN
		enlarged_shape.height = piece_shape.height + OVERLAP_MARGIN
	else:
		return null
	
	var space_state = builder_world.direct_space_state
	var query = PhysicsShapeQueryParameters2D.new()
	query.set_shape(enlarged_shape)
	query.transform = Transform2D(0, position)
	
	var result = space_state.intersect_shape(query, 1)
	
	if not result.is_empty():
		var collider = result[0].get("collider")
		if collider is StructuralPiece:
			# --- REFACTORED: The module is now the direct parent/owner ---
			if is_instance_valid(collider.owner) and collider.owner is Module:
				return collider.owner
				
	return null

func _find_first_module() -> Module:
	for node in builder_scene_root.get_children():
		if node is Module:
			return node
	return null

func _remove_piece_under_mouse():
	var viewport: SubViewport = main_screen.find_child("SubViewport")
	if not viewport: return
	var world_mouse_pos = viewport.get_canvas_transform().affine_inverse() * viewport.get_mouse_position()
	
	var space_state = builder_world.direct_space_state
	var query = PhysicsPointQueryParameters2D.new()
	query.position = world_mouse_pos
	var result = space_state.intersect_point(query, 1)

	if not result.is_empty():
		var collider = result[0].get("collider")
		if collider is StructuralPiece:
			_remove_piece_with_undo_redo(collider)
		elif collider is Component:
			pass


# --- REFACTORED: Piece Removal ---
func _remove_piece_with_undo_redo(piece: StructuralPiece):
	var module = piece.owner as Module
	if not is_instance_valid(module) or not module is Module:
		return
	
	undo_redo.create_action("Remove Structural Piece")
	
	# If this is the last structural piece of the module...
	if module.get_structural_pieces().size() == 1:
		# ...remove the entire module.
		undo_redo.add_do_method(builder_scene_root, "remove_child", module)
		undo_redo.add_undo_method(builder_scene_root, "add_child", module)
		undo_redo.add_undo_method(module, "set_owner", builder_scene_root) 
	else:
		# Otherwise, just remove the single piece from its parent (the module).
		undo_redo.add_do_method(module, "remove_child", piece)
		undo_redo.add_do_method(module, "_recalculate_collision_shape")

		undo_redo.add_undo_method(module, "add_child", piece)
		undo_redo.add_undo_method(piece, "set_owner", module) # Re-assign owner on undo
		undo_redo.add_undo_method(module, "_recalculate_collision_shape")

	undo_redo.commit_action()

# --- Toolbar Button Functions (No changes needed) ---
func _on_rotate_button_pressed():
	rotation_angle = wrapf(rotation_angle + PI / 2, 0, TAU)
	if is_instance_valid(preview_node):
		preview_node.rotation = rotation_angle
		_update_preview_position()

func _on_center_button_pressed():
	builder_camera.position = Vector2.ZERO
	builder_camera.zoom = Vector2(1.0, 1.0)
	
func _on_pressurise_button_pressed():
	pass

func _on_save_button_pressed():
	var module_to_save: Module
	var selected_nodes = EditorInterface.get_selection().get_selected_nodes()
	if not selected_nodes.is_empty() and selected_nodes[0] is Module:
		module_to_save = selected_nodes[0]
	else:
		module_to_save = _find_first_module()

	if not is_instance_valid(module_to_save):
		push_error("Error: No Module node found or selected to save.")
		return
	
	var save_dialog = EditorFileDialog.new()
	save_dialog.file_mode = EditorFileDialog.FILE_MODE_SAVE_FILE
	save_dialog.add_filter("*.tscn; Godot Scene")
	save_dialog.current_path = "res://modules/" + module_to_save.name + ".tscn"
	
	EditorInterface.get_editor_main_screen().add_child(save_dialog)
	save_dialog.popup_centered_ratio()

	save_dialog.file_selected.connect(Callable(self, "_perform_save").bind(module_to_save))

func _perform_save(file_path: String, module_to_save: Module):
	var save_dir = file_path.get_base_dir()
	var dir = DirAccess.open("res://")
	if not dir.dir_exists(save_dir):
		dir.make_dir_recursive(save_dir)
	
	var packed_scene = PackedScene.new()
	var error = packed_scene.pack(module_to_save)
	
	if error != OK:
		push_error("Error packing scene: ", error_string(error))
		return
	
	var save_result = ResourceSaver.save(packed_scene, file_path)
	
	if save_result == OK:
		print("Module saved successfully to ", file_path)
	else:
		push_error("Error saving scene: ", error_string(save_result))
		
	EditorInterface.get_resource_filesystem().scan()

func _on_undo_redo_action_committed():
	_refresh_tree_display()

# --- REFACTORED: Tree Display ---
func _refresh_tree_display():
	if not is_instance_valid(tree_control):
		return
		
	tree_control.clear()
	var root_item = tree_control.create_item()
	root_item.set_text(0, builder_scene_root.name)
	
	# Iterate through all modules and populate the tree.
	for module in builder_scene_root.get_children():
		if module is Module:
			var module_item = tree_control.create_item(root_item)
			module_item.set_text(0, module.name)
			module_item.set_meta("node", module)
			
			# Use the module's helper functions to find children
			for piece in module.get_structural_pieces():
				var piece_item = tree_control.create_item(module_item)
				piece_item.set_text(0, piece.name)
				piece_item.set_meta("node", piece)

			for component in module.get_components():
				var component_item = tree_control.create_item(module_item)
				component_item.set_text(0, component.name)
				component_item.set_meta("node", component)


func _find_closest_attachment_point(module: Module, world_pos: Vector2):
	var min_distance_sq = module.COMPONENT_GRID_SIZE * module.COMPONENT_GRID_SIZE * 0.5 
	var closest_point = null

	for point in module.get_attachment_points():
		var dist_sq = point.position.distance_squared_to(world_pos)
		if dist_sq < min_distance_sq:
			min_distance_sq = dist_sq
			closest_point = point
			
	return closest_point

const GeneratorScript = preload("res://data/structure/structure_generator.gd")

# The callback function
func _on_generate_structures_pressed():
	if GeneratorScript:
		var generator = GeneratorScript.new()
		if generator.has_method("generate_system_one"):
			generator.generate_system_one()
		else:
			push_error("StructureGenerator script missing 'generate_system_one' method.")

		if generator.has_method("generate_system_two_pentagonal"):
			generator.generate_system_two_pentagonal()
		else:
			push_error("StructureGenerator script missing 'generate_system_two_pentagonal' method.")
		
		if generator.has_method("generate_system_two_v2_sphere"):
			generator.generate_system_two_v2_sphere()
		else:
			push_error("StructureGenerator script missing 'generate_system_two_v2_sphere' method.")

		# Cleanup if it's a Node
		if generator is Node:
			generator.queue_free()