# encoding: utf-8

require 'one_apm/collector/sampler'
require 'one_apm/support/vm'

module OneApm
  module Collector
    module Samplers
      class VMSampler < Sampler
        OA_GC_RUNS_METRIC                = 'RubyVM/GC/runs'.freeze
        OA_HEAP_LIVE_METRIC              = 'RubyVM/GC/heap_live'.freeze
        OA_HEAP_FREE_METRIC              = 'RubyVM/GC/heap_free'.freeze
        OA_THREAD_COUNT_METRIC           = 'RubyVM/Threads/all'.freeze
        OA_OBJECT_ALLOCATIONS_METRIC     = 'RubyVM/GC/total_allocated_object'.freeze
        OA_MAJOR_GC_METRIC               = 'RubyVM/GC/major_gc_count'.freeze
        OA_MINOR_GC_METRIC               = 'RubyVM/GC/minor_gc_count'.freeze
        OA_METHOD_INVALIDATIONS_METRIC   = 'RubyVM/CacheInvalidations/method'.freeze
        OA_CONSTANT_INVALIDATIONS_METRIC = 'RubyVM/CacheInvalidations/constant'.freeze

        attr_reader :transaction_count

        named :vm

        def initialize
          @lock = Mutex.new
          @transaction_count = 0
          @last_snapshot = take_snapshot
        end

        def take_snapshot
          OneApm::Support::VM.snapshot
        end

        def setup_events(event_listener)
          event_listener.subscribe(:transaction_finished, &method(:on_transaction_finished))
        end

        def on_transaction_finished(*_)
          @lock.synchronize { @transaction_count += 1 }
        end

        def reset_transaction_count
          @lock.synchronize do
            old_count = @transaction_count
            @transaction_count = 0
            old_count
          end
        end

        def record_gc_runs_metric(snapshot, txn_count)
          if snapshot.gc_total_time || snapshot.gc_runs
            if snapshot.gc_total_time
              gc_time = snapshot.gc_total_time - @last_snapshot.gc_total_time.to_f
            end
            if snapshot.gc_runs
              gc_runs = snapshot.gc_runs - @last_snapshot.gc_runs
            end
            wall_clock_time = snapshot.taken_at - @last_snapshot.taken_at
            OneApm::Manager.agent.stats_engine.tl_record_unscoped_metrics(OA_GC_RUNS_METRIC) do |stats|
              stats.call_count           += txn_count
              stats.total_call_time      += gc_runs if gc_runs
              stats.total_exclusive_time += gc_time if gc_time
              stats.max_call_time         = (gc_time.nil? ? 0 : 1)
              stats.sum_of_squares       += wall_clock_time
            end
          end
        end

        def record_delta(snapshot, key, metric, txn_count)
          value = snapshot.send(key)
          if value
            delta = value - @last_snapshot.send(key)
            OneApm::Manager.agent.stats_engine.tl_record_unscoped_metrics(metric) do |stats|
              stats.call_count      += txn_count
              stats.total_call_time += delta
            end
          end
        end

        def record_gauge_metric(metric_name, value)
          OneApm::Manager.agent.stats_engine.tl_record_unscoped_metrics(metric_name) do |stats|
            stats.call_count      = value
            stats.sum_of_squares  = 1
          end
        end

        def record_heap_live_metric(snapshot)
          if snapshot.heap_live
            record_gauge_metric(OA_HEAP_LIVE_METRIC, snapshot.heap_live)
          end
        end

        def record_heap_free_metric(snapshot)
          if snapshot.heap_free
            record_gauge_metric(OA_HEAP_FREE_METRIC, snapshot.heap_free)
          end
        end

        def record_thread_count_metric(snapshot)
          if snapshot.thread_count
            record_gauge_metric(OA_THREAD_COUNT_METRIC, snapshot.thread_count)
          end
        end

        def poll
          snap = take_snapshot
          tcount = reset_transaction_count

          record_gc_runs_metric(snap, tcount)
          record_delta(snap, :total_allocated_object, OA_OBJECT_ALLOCATIONS_METRIC, tcount)
          record_delta(snap, :major_gc_count, OA_MAJOR_GC_METRIC, tcount)
          record_delta(snap, :minor_gc_count, OA_MINOR_GC_METRIC, tcount)
          record_delta(snap, :method_cache_invalidations, OA_METHOD_INVALIDATIONS_METRIC, tcount)
          record_delta(snap, :constant_cache_invalidations, OA_CONSTANT_INVALIDATIONS_METRIC, tcount)
          record_heap_live_metric(snap)
          record_heap_free_metric(snap)
          record_thread_count_metric(snap)

          @last_snapshot = snap
        end
      end
    end
  end
end