require 'prometheus/client'
require 'mmap'

module Prometheus
  module Client
    module Helper
      class MmapedFile < Mmap
        class << self
          def open_readonly(filepath)
            MmapedFile.new(filepath, 'r', Mmap::MAP_PRIVATE)
          end

          def open(filepath)
            MmapedFile.new(filepath, 'rw', Mmap::MAP_SHARED)
          end

          def lock_to_process(filepath)
            m = open(filepath)
            if m.lock_to_process
              Kernel.at_exit do
                open(filepath).unlock
              end
              true
            else
              false
            end
          ensure
            m.close
          end

          def ensure_process_exclusive_file(file_prefix = 'mmaped_file')
            (0..Float::INFINITY).lazy
              .map { |f_num| "#{file_prefix}_#{Prometheus::Client.pid}-#{f_num}.db" }
              .map { |filename| File.join(Prometheus::Client.configuration.multiprocess_files_dir, filename) }
              .find { |path| Helper::MmapedFile.lock_to_process(path) }
          end

          def open_process_exclusive_file(file_prefix = 'mmaped_file')
            filename = Helper::MmapedFile.ensure_process_exclusive_file(file_prefix)
            open(filename)
          end
        end

        MINIMUM_SIZE = 8
        attr_reader :filepath, :size

        def initialize(filepath, mode = 'r', protection = Mmap::MAP_SHARED, options = {})
          @filepath = filepath

          File.open(filepath, 'a+b') do |file|
            file.truncate(initial_mmap_file_size) if file.size < MINIMUM_SIZE
            @size = file.size
          end

          super(filepath, mode, protection, options)
        end

        def pid
          self[0..3].unpack('l')[0]
        end

        def pid=(value)
          self[0..3] = [value].pack('l')
        end

        def used
          self[4..7].unpack('l')[0]
        end

        def used=(value)
          self[4..7] = [value].pack('l')
        end

        def locked_to_process?
          pid > 0 && pid_alive?(pid)
        end

        def lock_owner?
          pid == Process.pid
        end

        def lock_to_process
          return true if lock_owner?

          if locked_to_process?
            false
          else
            self.pid = Process.pid

            # check if PID was correctly written
            lock_owner?
          end
        end

        def unlock
          return unless lock_owner?

          self.pid = 0
        end

        def entries
          return Enumerator.new {} if used.zero?

          Enumerator.new do |yielder|
            used_ = used # cache used to avoid unnecessary unpack operations

            pos = 12 # pid + used + padding offset
            while pos < used_
              data = slice(pos..-1)
              encoded_len, = data.unpack('l')
              padding_len = 8 - (encoded_len + 4) % 8
              value_offset = 4 + encoded_len + padding_len
              pos += value_offset

              yielder.yield data, encoded_len, value_offset, pos

              pos += 8
            end
          end
        end

        def add_entry(data, value)
          self.used = 12 if used.zero?

          # Pad to be 8-byte aligned.
          padded = data + (' ' * (8 - (data.length + 4) % 8))
          entry = [data.length, padded, value].pack("lA#{padded.length}d")
          used_ = used

          while (used_ + entry.length) > @size do
            extend(size)
            @size = File.size(filepath)
          end

          self[used_..used_ + entry.length] = entry

          self.used = used_ + entry.length
        end

        def close
          munmap
        end

        private

        def initial_mmap_file_size
          Prometheus::Client.configuration.initial_mmap_file_size
        end

        def pid_alive?(pid)
          begin
            Process.getpgid(pid)
            true
          rescue Errno::ESRCH
            false
          end
        end
      end
    end
  end
end