# frozen_string_literal: true

# Released under the MIT License.
# Copyright, 2025-2026, by Samuel Williams.

require "async/bus/controller"
require "stringio"

module Async
	module Service
		module Supervisor
			# Controller for worker operations.
			#
			# Handles diagnostic operations like memory dumps, thread dumps, scheduler dumps, etc.
			class WorkerController < Async::Bus::Controller
				# Initialize the worker controller for the worker.
				def initialize(worker)
					@worker = worker
				end
				
				private def dump(path: nil, log: nil, buffer: true, &block)
					if path
						File.open(path, "w") do |file|
							yield file
						end
						
						return {path: path}
					elsif buffer
						buffer = StringIO.new
						yield buffer
						
						if log
							Console.info(self, log, data: buffer.string)
							return {}
						else
							return {data: buffer.string}
						end
					else
						raise ArgumentError, "Buffered output not supported!"
					end
				end
				
				# Dump the current fiber scheduler hierarchy.
				#
				# Generates a hierarchical view of all running fibers and their relationships.
				#
				# @parameter path [String] Optional file path to save the dump.
				# @parameter log [String] Optional log message to output.
				def scheduler_dump(path: nil, log: nil)
					dump(path: path, log: log) do |file|
						Fiber.scheduler.print_hierarchy(file)
					end
				end
				
				# Dump the entire object space to a file.
				#
				# This is a heavyweight operation that dumps all objects in the heap.
				#
				# @parameter path [String] Optional file path to save the dump.
				# @parameter shapes [Boolean] Whether to include Ruby shape-tree records.
				def memory_dump(path: nil, shapes: true)
					require "objspace"
					
					dump(path: path, buffer: false) do |file|
						ObjectSpace.dump_all(output: file, shapes: shapes)
					end
				end
				
				# Start recording object allocation metadata.
				#
				# Allocation tracing is process-global and can add significant overhead.
				def allocation_trace_start
					require "objspace"
					
					raise "Object allocation tracing was already started by this controller!" if @allocation_trace_active
					raise "Object allocation tracing is already active!" if ObjectSpace.allocation_sourcefile(Object.new)
					
					ObjectSpace.trace_object_allocations_start
					@allocation_trace_active = true
					
					return {started: true}
				end
				
				# Stop recording allocations, dump the traced heap, and clear the metadata.
				#
				# @parameter path [String] File path where the heap dump should be written.
				# @parameter shapes [Boolean] Whether to include Ruby shape-tree records.
				def allocation_trace_stop(path:, shapes: true)
					require "objspace"
					stopped = false
					
					raise "Object allocation tracing was not started by this controller!" unless @allocation_trace_active
					
					ObjectSpace.trace_object_allocations_stop
					stopped = true
					@allocation_trace_active = false
					
					memory_dump(path: path, shapes: shapes)
				ensure
					ObjectSpace.trace_object_allocations_clear if stopped
				end
				
				# Dump information about all running threads.
				#
				# Includes thread inspection and backtraces for debugging.
				#
				# @parameter path [String] Optional file path to save the dump.
				def thread_dump(path: nil)
					dump(path: path) do |file|
						Thread.list.each do |thread|
							file.puts(thread.inspect)
							file.puts(thread.backtrace)
						end
					end
				end
				
				# Start garbage collection profiling.
				#
				# Enables the GC profiler to track garbage collection performance.
				#
				# @returns [Hash] Confirmation that profiling started.
				def garbage_profile_start
					GC::Profiler.enable
					return {started: true}
				end
				
				# Stop garbage collection profiling and return results.
				#
				# Disables the GC profiler and returns collected profiling data.
				#
				# @parameter path [String] Optional file path to save the profile.
				def garbage_profile_stop(path: nil)
					dump(path: path) do |file|
						file.puts GC::Profiler.result
					end
				ensure
					GC::Profiler.disable
				end
				
				# Setup utilization observer for this worker.
				#
				# Delegates to the worker to map the shared memory file and configure its
				# utilization registry. Called by the supervisor via RPC.
				#
				# @parameter path [String] Path to the shared memory file that the worker should map.
				# @parameter size [Integer] Size of the shared memory region to map.
				# @parameter offset [Integer] Offset into the shared memory buffer allocated for this worker.
				# @returns [Array] Array of [key, type, offset] tuples describing the utilization schema.
				def setup_utilization_observer(path, size, offset)
					@worker.setup_utilization_observer(path, size, offset)
				end
			end
		end
	end
end
