# encoding: utf-8

require 'one_apm/collector/commands/agent_command'
require 'one_apm/collector/commands/xray_session_collection'
require 'one_apm/collector/commands/thread_profiler_session'
require 'one_apm/collector/commands/restart_agent'
require 'one_apm/support/backtrace/backtrace_service'

module OneApm
  module Collector
    class AgentCommandRouter
      attr_reader :handlers

      attr_accessor :thread_profiler_session,
                    :backtrace_service,
                    :xray_session_collection

      def initialize(event_listener=nil)
        @handlers    = Hash.new { |*| Proc.new { |cmd| self.unrecognized_agent_command(cmd) } }

        @backtrace_service = OneApm::Agent::Threading::BacktraceService.new(event_listener)

        @thread_profiler_session = OneApm::Collector::Commands::ThreadProfilerSession.new(@backtrace_service)
        @xray_session_collection = OneApm::Collector::Commands::XraySessionCollection.new(@backtrace_service, event_listener)

        @handlers['restart'] = Proc.new { OneApm::Collector::Commands::RestartAgent.new }
        @handlers['start_profiler'] = Proc.new { |cmd| thread_profiler_session.handle_start_command(cmd) }
        @handlers['stop_profiler']  = Proc.new { |cmd| thread_profiler_session.handle_stop_command(cmd) }
        @handlers['active_xray_sessions'] = Proc.new { |cmd| xray_session_collection.handle_active_xray_sessions(cmd) }

        if event_listener
          event_listener.subscribe(:before_shutdown, &method(:on_before_shutdown))
        end
      end

      def one_apm_service
        OneApm::Manager.agent.service
      end

      def check_for_and_handle_agent_commands
        commands = get_agent_commands

        stop_xray_sessions unless active_xray_command?(commands)

        results = invoke_commands(commands)
        one_apm_service.agent_command_results(results) unless results.empty?
      end

      def stop_xray_sessions
        self.xray_session_collection.stop_all_sessions
      end

      def active_xray_command?(commands)
        commands.any? {|command| command.name == 'active_xray_sessions'}
      end

      def on_before_shutdown(*args)
        if self.thread_profiler_session.running?
          self.thread_profiler_session.stop(true)
        end
      end

      def harvest!
        profiles = []
        profiles += harvest_from_xray_session_collection
        profiles += harvest_from_thread_profiler_session
        log_profiles(profiles)
        profiles
      end

      # We don't currently support merging thread profiles that failed to send
      # back into the AgentCommandRouter, so we just no-op this method.
      # Same with reset! - we don't support asynchronous cancellation of a
      # running thread profile or X-Ray session currently.
      def merge!(*args); end
      def reset!; end

      def harvest_from_xray_session_collection
        self.xray_session_collection.harvest_thread_profiles
      end

      def harvest_from_thread_profiler_session
        if self.thread_profiler_session.ready_to_harvest?
          self.thread_profiler_session.stop(true)
          [self.thread_profiler_session.harvest]
        else
          []
        end
      end

      def log_profiles(profiles)
        if profiles.empty?
          OneApm::Manager.logger.debug "No thread profiles with data found to send."
        else
          profile_descriptions = profiles.map { |p| p.to_log_description }
          OneApm::Manager.logger.debug "Sending thread profiles [#{profile_descriptions.join(", ")}]"
        end
      end

      def get_agent_commands
        commands = one_apm_service.get_agent_commands || []
        OneApm::Manager.logger.info "Received get_agent_commands = #{commands.inspect}" if commands.any?
        commands.map {|collector_command| OneApm::Collector::Commands::AgentCommand.new(collector_command)}
      end

      def invoke_commands(agent_commands)
        results = {}

        agent_commands.each do |agent_command|
          results[agent_command.id.to_s] = invoke_command(agent_command)
        end

        results
      end

      class AgentCommandError < StandardError
      end

      def invoke_command(agent_command)
        begin
          call_handler_for(agent_command)
          return success
        rescue AgentCommandError => e
          OneApm::Manager.logger.debug(e)
          error(e)
        end
      end

      OA_SUCCESS_RESULT = {}.freeze
      OA_ERROR_KEY = "error"

      def success
        OA_SUCCESS_RESULT
      end

      def error(err)
        { OA_ERROR_KEY => err.message }
      end

      def call_handler_for(agent_command)
        handler = select_handler(agent_command)
        handler.call(agent_command)
      end

      def select_handler(agent_command)
        @handlers[agent_command.name]
      end

      def unrecognized_agent_command(agent_command)
        OneApm::Manager.logger.debug("Unrecognized agent command #{agent_command.inspect}")
      end
    end
  end
end