# 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/telemetry'

module Contrast
  module Agent
    # This class used to ensure that our worker threads are running in multi-process environments
    #
    # @attr_reader heapdump_util [Contrast::Utils::HeapDumpUtil]
    # @attr_reader heartbeat [Contrast::Agent::ServiceHeartbeat]
    # @attr_reader messaging_queue [Contrast::Api::Communication::MessagingQueue]
    class ThreadWatcher
      include Contrast::Components::Logger::InstanceMethods

      attr_reader :heapdump_util, :heartbeat, :messaging_queue

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

      def startup!
        return unless ::Contrast::AGENT.enabled?

        telemetry_status = telemetry_thread_init
        heartbeat_status = heartbeat_thread_init
        messaging_status = messaging_thread_init
        reporter_status = reporter_thread_init

        logger.debug('ThreadWatcher started threads')

        @pids[Process.pid] = messaging_status && heartbeat_status
        return @pids unless Contrast::Agent::Telemetry.enabled?

        @pids[Process.pid] = @pids[Process.pid] && telemetry_status
        return @pids unless Contrast::Agent::Reporter.enabled?

        @pids[Process.pid] = @pids[Process.pid] && reporter_status
      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!
      end

      def heartbeat_thread_init
        unless heartbeat.running?
          logger.debug('Attempting to start heartbeat thread')
          heartbeat.start_thread!
        end
        heartbeat_result = heartbeat.running?
        logger.debug('Heartbeat thread status', alive: heartbeat_result)
        heartbeat_result
      end

      def telemetry_thread_init
        @telemetry.start_thread! if @telemetry&.attempt_to_start?
        telemetry_result = @telemetry&.running?
        logger.debug('Telemetry thread status', alive: telemetry_result)
        telemetry_result
      end

      def messaging_thread_init
        unless messaging_queue.running?
          logger.debug('Attempting to start messaging queue thread')
          messaging_queue.start_thread!
        end
        messaging_result = messaging_queue.running?
        logger.debug('Messaging thread status', alive: messaging_result)
        messaging_result
      end

      def reporter_thread_init
        @reporter.start_thread! if @reporter&.attempt_to_start?
        unless @reporter&.running?
          logger.debug('Attempting to start reporter thread')
          @reporter&.start_thread!
        end
        reporter_result = @reporter&.running?
        logger.debug('Reporter thread status', alive: reporter_result)
        reporter_result
      end

      # @return [Contrast::Agent::Telemetry]
      def telemetry_queue
        return if @telemetry.nil?

        @telemetry
      end

      # @return [Contrast::Agent::Reporter]
      def reporter
        return if @reporter.nil?

        @reporter
      end
    end
  end
end