#!/usr/bin/env ruby

require 'pathname'
require 'prometheus/client'
require 'lxc'
require 'ostruct'
require 'pmap'

require_relative 'ns'
require_relative 'file-statfs'
require_relative 'cborio'
require_relative 'lxc_extend'
require_relative 'lxc-exporter/os_release'
require_relative 'lxc-exporter/mount_info'

class LxcCollector
  include Prometheus::Client

  CTStates = { aborting: -1, stopped: 0, starting: 0.75, stopping: 0.25, running: 1 }

  def statfs ct, mis
    mis.each do |mi|
      next if %w[devpts devtmpfs fuse.lxcfs].include? mi.type
      statfs = File::Statfs.new mi.mp
      yield ct, mi.mp, statfs
    end
  end

  def initialize registry = nil
    @prometheus ||= Registry.new
    @up     = Gauge.new :lxc_up, labels: %i[id name], docstring: 'Is CT running?'
    @state  = Gauge.new :lxc_state, labels: %i[id name],
      docstring: "State of CT (#{CTStates.map{|n,i|"#{i}=#{n}"}.join', '})"
    @info   = Gauge.new :lxc_filesystem_info,           
      labels: %i[id mountpoint name fstype], docstring: 'CT Info'
    @totalb = Gauge.new :lxc_filesystem_total_bytes,    
      labels: %i[id mountpoint], docstring: 'Total bytes in filesystem'
    @freeb  = Gauge.new :lxc_filesystem_free_bytes,     
      labels: %i[id mountpoint], docstring: 'Free bytes in filesystem'
    @availb = Gauge.new :lxc_filesystem_available_bytes,
      labels: %i[id mountpoint], docstring: 'Free bytes available to unprivileged user'
    @totalf = Gauge.new :lxc_filesystem_total_files,    
      labels: %i[id mountpoint], docstring: 'Total files in filesystem'
    @freef  = Gauge.new :lxc_filesystem_free_files,     
      labels: %i[id mountpoint], docstring: 'Free files in filesystem'
    @os_release = Gauge.new :lxc_os_release_info,       
      labels: %i[id name os_id os_name os_pretty_name os_version_codename],
      docstring: 'OS release info'
    @os_version = Gauge.new :lxc_os_version_id,
      labels: %i[id], docstring: 'OS release (distribution) version'
    @pkgs_upgradable = Gauge.new :lxc_packages_upgradable,
      labels: %i[id], docstring: 'Count of upgradable packages'
    @pkgs_last_update = Gauge.new :lxc_packages_last_update,
      labels: %i[id], docstring: 'Last successfull source update'
    [
      @up, @state,
      @info, @totalb, @freeb, @availb, @totalf, @freef,
      @os_release, @os_version,
      @pkgs_upgradable, @pkgs_last_update
    ].each {|g| @prometheus.register g }
  end

  AptLastUpdateFile = Pathname.new '/var/cache/apt/pkgcache.bin'

  def forkedo wr
    LXC.running_containers do |ct|
      pid = ct.init_pid
      mis = MountInfo.of pid
      NS.change pid, :pid, :mnt do
        name = ct.config_item( 'lxc.uts.name')

        # Filesystems
        statfs ct, mis do |ct, mp, st|
          next  unless st.fstypename and 0 < st.blocks
          st = st.to_h.merge mp: mp, id: ct.name, name: name
          wr.write [:mountpoint, st]
        end

        # OS / Distribution
        osr = OsRelease.read
        if osr
          wr.write [:os_release, {
            id: ct.name, name: name,
            os_id:               osr[:id],
            os_name:             osr[:name],
            os_pretty_name:      osr[:pretty_name],
            os_version_id:       osr[:version_id].to_i,
            os_version_codename: osr[:version_codename]
          }]

          case osr[:id].to_sym
          # Special things for Debian(-based) distributions
          when :debian
            upgradable = nil
            pid = 
              IO.popen %w[apt list --upgradable], err: "/dev/null" do |apt|
                upgradable = apt.readlines.length - 1
              end
            wr.write [:pkgs, {
              id: ct.name, name: name, upgradable: upgradable,
              last_update: AptLastUpdateFile.exist? ? AptLastUpdateFile.stat.mtime.to_f : -1,
            }]
          end
        end
      end
    end
    wr.close
    exit 0
  end

  def collecto
    LXC.containers do |ct|
      labels = { id: ct.name, name: ct.config_item( 'lxc.uts.name') }
      @up.set ct.running? ? 1 : 0, labels: labels
      @state.set CTStates[ct.state], labels: labels
    end
    pid, rd = CBORIO.popen( &method( :forked))
    rd.each do |l|
      case l[0].to_sym

      when :mountpoint
        st = OpenStruct.new l[1]
        labels = { id: st.id, mountpoint: st.mp }
        @info.set 1, labels: labels.merge( name: st.name, fstype: st.fstypename)
        @totalb.set st.bsize*st.blocks, labels: labels
        @freeb.set  st.bsize*st.bfree,  labels: labels
        @availb.set st.bsize*st.bavail, labels: labels
        @totalf.set st.files, labels: labels
        @freef.set  st.ffree, labels: labels

      when :os_release
        osr = OpenStruct.new l[1]
        @os_release.set 1, labels: {
          id: osr.id, name: osr.name,
          os_id: osr.os_id,
          os_name: osr.os_name,
          os_pretty_name: osr.os_pretty_name,
          os_version_codename: osr.os_version_codename
        }
        @os_version.set osr.os_version_id, labels: { id: osr.id }

      when :pkgs
        as = OpenStruct.new l[1]
        @pkgs_last_update.set as.last_update, labels: { id: as.id }
        @pkgs_upgradable.set  as.upgradable,  labels: { id: as.id }
      end
    end
    rd.close
    Process.wait pid
    @prometheus
  end

  def forked2 io, ct
    pid = ct.init_pid
    mis = MountInfo.of pid
    id, name = ct.name, ct.config_item( 'lxc.uts.name')
    NS.enter pid, keep_open: [io, STDERR, STDOUT]

    # Filesystems
    statfs ct, mis do |ct, mp, st|
      next  unless st.fstypename and 0 < st.blocks
      st = st.to_h.merge mp: mp, id: id, name: name
      io.write [:mountpoint, st]
    end

    # OS / Distribution
    osr = OsRelease.read
    unless osr.nil?
      io.write [:os_release, {
        id: id, name: name,
        os_id:               osr[:id],
        os_name:             osr[:name],
        os_pretty_name:      osr[:pretty_name],
        os_version_id:       osr[:version_id].to_i,
        os_version_codename: osr[:version_codename]
      }]

      case osr[:id].to_sym
      # Special things for Debian(-based) distributions
      when :debian
        upgradable = nil
        pid = 
          IO.popen %w[apt list --upgradable], err: "/dev/null" do |apt|
            upgradable = apt.readlines.length - 1
          end
        io.write [:pkgs, {
          id: id, name: name, upgradable: upgradable,
          last_update: AptLastUpdateFile.stat.mtime.to_f,
        }]
      end
    end
  rescue Object
    STDERR.puts "#{$!} (#{$!.class})", $!.backtrace.map {|x| "  #{x}"}
  end

  def collect2
    LXC.containers do |ct|
      labels = { id: ct.name, name: ct.config_item( 'lxc.uts.name') }
      @up.set ct.running? ? 1 : 0, labels: labels
      @state.set CTStates[ct.state], labels: labels
    end

    LXC.running_containers.peach do |ct|
      pid, io = CBORIO.popen(:r) {|io| forked2 io, ct }
      begin
        io.map do |l|
          case l[0].to_sym

          when :mountpoint
            st = OpenStruct.new l[1]
            labels = { id: st.id, mountpoint: st.mp }
            @info.set 1, labels: labels.merge( name: st.name, fstype: st.fstypename)
            @totalb.set st.bsize*st.blocks, labels: labels
            @freeb.set  st.bsize*st.bfree,  labels: labels
            @availb.set st.bsize*st.bavail, labels: labels
            @totalf.set st.files, labels: labels
            @freef.set  st.ffree, labels: labels

          when :os_release
            osr = OpenStruct.new l[1]
            @os_release.set 1, labels: {
              id: osr.id, name: osr.name,
              os_id: osr.os_id,
              os_name: osr.os_name,
              os_pretty_name: osr.os_pretty_name,
              os_version_codename: osr.os_version_codename
            }
            @os_version.set osr.os_version_id, labels: { id: osr.id }

          when :pkgs
            as = OpenStruct.new l[1]
            @pkgs_last_update.set as.last_update, labels: { id: as.id }
            @pkgs_upgradable.set  as.upgradable,  labels: { id: as.id }
          end
        end
      ensure
        io.close
        Process.wait pid
      end
    end
    @prometheus
  end

  alias forked forked2
  alias collect collect2
end

if Pathname.new( __FILE__).expand_path == Pathname.new( $0).expand_path
  require 'prometheus/client/formats/text'
  puts Prometheus::Client::Formats::Text.marshal( LxcCollector.new.collect2)
end