# Copyright (c) 2022 Contrast Security, Inc. See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
# frozen_string_literal: true

require 'contrast/components/logger'
require 'contrast/agent/service_heartbeat'
require 'contrast/api/communication/messaging_queue'
require 'contrast/agent/reporting/report'
require 'contrast/agent/reporting/reporter_heartbeat'
require 'contrast/agent/telemetry/base'

module Contrast
  module Agent
    # This class used to ensure that our worker threads are running in multi-process environments
    class ThreadWatcher
      include Contrast::Components::Logger::InstanceMethods

      # @return [Contrast::Utils::HeapDumpUtil]
      attr_reader :heapdump_util
      # @return [Contrast::Agent::ServiceHeartbeat, nil]
      attr_reader :heartbeat
      # @return [Contrast::Api::Communication::MessagingQueue, nil]
      attr_reader :messaging_queue
      # @return [Contrast::Agent::Reporter]
      attr_reader :reporter
      # @return [Contrast::Agent::ReporterHeartbeat]
      attr_reader :reporter_heartbeat

      def initialize
        @pids = {}
        @heapdump_util = Contrast::Utils::HeapDumpUtil.new
        unless ::Contrast::CONTRAST_SERVICE.unnecessary?
          @heartbeat = Contrast::Agent::ServiceHeartbeat.new
          @messaging_queue = Contrast::Api::Communication::MessagingQueue.new
        end
        @reporter = Contrast::Agent::Reporter.new
        @reporter_heartbeat = Contrast::Agent::ReporterHeartbeat.new
        @telemetry = Contrast::Agent::Telemetry::Base.new if Contrast::Agent::Telemetry::Base.enabled?
      end

      # @return [Hash, nil] map of process to thread startup status
      def startup!
        return unless ::Contrast::AGENT.enabled?

        logger.debug('ThreadWatcher started threads')
        heartbeat_status = init_thread(heartbeat)
        messaging_status = init_thread(messaging_queue)
        @pids[Process.pid] = messaging_status && heartbeat_status
        if Contrast::Agent::Telemetry::Base.enabled?
          telemetry_status = init_thread(telemetry_queue)
          @pids[Process.pid] = @pids[Process.pid] && telemetry_status
        end
        reporter_status = init_thread(reporter)
        reporter_heartbeat_status = init_thread(reporter_heartbeat)
        @pids[Process.pid] = @pids[Process.pid] && reporter_status && reporter_heartbeat_status
        @pids
      end

      def ensure_running?
        return if @pids[Process.pid] == true

        logger.debug('ThreadWatcher - threads not running')
        startup!
      end

      def shutdown!
        heartbeat&.stop!
        messaging_queue&.stop!
        heapdump_util&.stop!
        telemetry_queue&.stop!
        reporter&.stop!
        reporter_heartbeat&.stop!
      end

      # @return [Contrast::Agent::Telemetry::Base]
      def telemetry_queue
        @telemetry
      end

      private

      # Start the thread governed by the given watcher
      #
      # @param watcher [Contrast::Agent::ThreadWatcher]
      # @return [Boolean] if the watched thread started successfully
      def init_thread watcher
        return unless watcher&.attempt_to_start?

        unless watcher.running?
          logger.debug('Attempting to start thread', type: watcher.to_s)
          watcher.start_thread!
        end
        result = watcher.running?
        logger.debug('Thread status', type: watcher.to_s, alive: result)
        result
      end
    end
  end
end