#!/bin/ruby
# encoding:utf-8
# frozen_string_literal: true

require 'rubygems'
require 'fileutils'
require 'json'
require 'logger'
require 'optparse'
require 'yaml'
require 'open3'

require_relative 'grafana_setup.rb'

class MultiIO
  def initialize(*targets)
    @targets = targets
  end

  def write(*args)
    @targets.each { |t| t.write(*args) }
  end

  def close
    @targets.each(&:close)
  end
end

module ScriptExporterSetUp
  PROP_ENABLED = 'wd.mon.pharos.script.exporter.enabled'.freeze
  GRAFANA_CONFIG_DIR = '/etc/grafana/conf.d/script_exporter_scrapes'.freeze
  SCRIPT_EXPORTER_CONFIG_DIR = '/etc/pharos_script_exporter/conf.d'.freeze
  SCRIPT_EXPORTER_CONFIG_FILE = '/etc/pharos_script_exporter/config.yaml'.freeze
  LOG_FILE = "/var/log/pharos_script_exporter/pharos_script_exporter_setup.log".freeze
  WPC_SYSTEM_PROPERTIES = '/etc/system.properties'.freeze
  DEFAULT_TIMEOUT = '55s'.freeze
  DEFAULT_INTERVAL = '60s'.freeze
  SCRIPT_TIMEOUT_ENFORCED = true

  @logger = Logger.new(STDOUT)
  @logger.level = Logger::INFO

  def self.logger
    @logger
  end

  class InvalidConfigError < StandardError; end

  def self.generate_grafana_agent_scrape(cfg)
    grafana_scrape = {
      'job_name' => cfg['name'],
      'metrics_path' => '/probe',
      'scrape_interval' => cfg.fetch('scrape_interval',DEFAULT_INTERVAL),
      'scrape_timeout' =>  cfg.fetch('scrape_timeout',DEFAULT_TIMEOUT),
      'add_source' => cfg.fetch('add_source', true),
      'params' => {
        'script' => [cfg['name']],
      },
      'static_configs' => [
        {
          'targets' => ['127.0.0.1:9469'],
          'labels' => cfg['labels']
        }
      ]
    }
    grafana_scrape['metric_relabel_configs'] = cfg.fetch('metric_relabel_configs', [])
    if cfg['prefix'] != nil
      grafana_scrape['metric_relabel_configs'] << {
        'source_labels' => ['__name__'],
        'target_label' => '__name__',
        'regex' => '.*',
        'replacement' => "#{cfg['prefix']}_${0}"
      }
    end
    grafana_scrape['params']['output'] = cfg['output'] != nil && !cfg['output'] ? ['ignore'] : []
    grafana_scrape['relabel_configs'] = cfg['relabel_configs'] if cfg['relabel_configs'] != nil
    return grafana_scrape
  end

  def self.generate_partial_script_exporter_config(cfg)
    script_exporter_cfg = {
      'name' => cfg['name'],
      'script' => cfg['command'],
      'timeout' => {
        'enforced'=> SCRIPT_TIMEOUT_ENFORCED
      }
    }
    script_exporter_cfg['timeout']['max_timeout'] = GrafanaConfigSeeder.in_secs(cfg.fetch('scrape_timeout',DEFAULT_TIMEOUT)) - 1
    return script_exporter_cfg
  end

  def self.generate_final_script_exporter_config(partial_script_exporter_cfgs)
    {
      'scripts' => partial_script_exporter_cfgs
    }
  end

  def self.validate_config(cfg)
    raise InvalidConfigError, "Missing name #{cfg.to_json}" unless cfg.key?('name')
    raise InvalidConfigError, "Missing script command #{cfg.to_json}" unless cfg.key?('command')
    raise InvalidConfigError, "Missing scrape labels #{cfg.to_json}" unless cfg.key?('labels')
    @logger.info "Config #{cfg['name']} is valid"
  end

  def self.config_changed(config_file_path, new_cfg)
    return true unless File.exist? config_file_path
    old_cfg = read_file_to_yaml(config_file_path)
    @logger.info "File #{config_file_path} changed: #{old_cfg != new_cfg}"
    old_cfg != new_cfg
  end

  def self.write_scrapes_to_file(grafana_agent_config_dir, scrapes)
    @logger.info "Writing grafana scrapes to #{grafana_agent_config_dir}"
    files_already_in_dir = Dir["#{grafana_agent_config_dir}/*.yaml"]
    scrape_names = []

    scrapes.each do |scrape|
      scrape_names << scrape["job_name"]

      file_name = File.join(grafana_agent_config_dir, "#{scrape["job_name"]}.yaml")

      if !files_already_in_dir.include?(file_name)
        write_config_to_file(file_name,scrape)
      elsif config_changed(file_name, scrape)
        @logger.info "Overwriting #{file_name} in #{grafana_agent_config_dir}"
        write_config_to_file(file_name,scrape)
      end
    end
    files_already_in_dir.each do |file|
      file_scrape_name = File.basename(file, ".yaml")

      if !scrape_names.include?(file_scrape_name)
        @logger.info "Removing #{file} as script exporter counterpart has been removed."
        File.delete(file) if File.exist?(file)
      end
    end
  end

  def self.read_file_to_yaml(file)
    begin
      unless File.exist? file
        @logger.warn "Cannot read file: #{file} as it does not exist"
        return []
      end
      YAML.load File.read file
    rescue => error
      @logger.warn "Cannot read file: #{file}: reason #{error} "
    end
  end

  def self.write_config_to_file(file, cfg)
    @logger.info "Attempting to write #{file} to file"
    begin
      File.open(file, 'w') do |f|
        f.write(cfg.to_yaml)
        @logger.info "Successfully wrote #{file} to file"
        f.chmod(0644)
      end
    rescue => error
      @logger.warn "Cannot write file: #{file}: reason #{error} "
    end
  end

  def self.load_config_properties(props_file)
    props = {}
    # Comment lines in .properties files are denoted by
    # the number sign (#) or the exclamation mark (!) as the first non blank character
    File.readlines(props_file)
        .reject { |line| line.lstrip[/^[#!]/] }
        .select { |line| line[/=/] }
        .each do |line|
      key, val = line.split('=', 2)
      props[key.strip] = val.strip
    end
    props
  end

  def self.run_command(cmd)
    @logger.info "Running command: #{cmd}"
    stdout, stderr, exit_status = Open3.capture3(cmd)
    @logger.info "#{cmd}: has succeeded - #{stdout} #{stderr}" if exit_status.success?
    @logger.error "#{cmd}: has failed - #{stdout} #{stderr}" unless exit_status.success?
    return exit_status
  end

  def self.main
    @logger = Logger.new MultiIO.new(STDOUT, File.open(LOG_FILE, "a"))
    @logger.level = Logger::INFO

    options = {}
    OptionParser.new do |opts|
      opts.banner = "Usage: pharos_script_exporter.rb [options]"

      opts.on("-c", "--check-enabled", "Starts the script exporter if it is enabled on this host otherwise stops it ") do |e|
        options[:check_enabled] = e
      end

      options[:script_exporter_config] = SCRIPT_EXPORTER_CONFIG_FILE
      opts.on("-c", "--script-exporter-config FILE", "the path to the script exporter config file, defaults to #{SCRIPT_EXPORTER_CONFIG_FILE}") do |file|
        options[:script_exporter_config] = file
      end

      options[:script_exporter_config_dir] = SCRIPT_EXPORTER_CONFIG_DIR
      opts.on("-d", "--script-exporter-conf-dir FILE", "the directory the script exporter will read configs from, defaults to #{SCRIPT_EXPORTER_CONFIG_DIR}") do |file|
        options[:script_exporter_config_dir] = file
      end

      options[:grafana_scrape_dir] = GRAFANA_CONFIG_DIR
      opts.on("-s", "--grafana-scrape-dir FOLDER",
              "directory containing extra scrapes yaml config, defaults to #{GRAFANA_CONFIG_DIR}.") do |folder|
        options[:grafana_scrape_dir] = folder
      end

      options[:properties_file] = WPC_SYSTEM_PROPERTIES
      opts.on("-p", "--properties-file FILE", "Custom system properties file, defaults to #{WPC_SYSTEM_PROPERTIES}") do |file|
        options[:properties_file] = file
      end

      opts.on('-h', '--help', 'Display this screen') do
        puts opts
        exit 0
      end

    end.parse!

    system_properties = load_config_properties(options[:properties_file])
    enabled = system_properties.fetch(PROP_ENABLED, 'false') == 'true'
    if options[:check_enabled]
      # only check if script exporter is enabled and exit
      if enabled
        @logger.info 'The Pharos Script Exporter is enabled, attempting to start agent service.'
        exit_status = run_command('service pharos_script_exporter start')
        return exit_status
      else
        @logger.info 'The Pharos Script Exporter is not enabled, attempting to stop agent service.'
        exit_status = run_command('service pharos_script_exporter stop')
        return exit_status
      end
    end

    unless enabled
      @logger.info 'The Pharos Script Exporter is not enabled, attempting to stop agent service.'
      run_command('service pharos_script_exporter stop')
      exit 1
    end

    unless Dir.exist? options[:script_exporter_config_dir]
      @logger.error "Missing configuration directory: #{options[:script_exporter_config_dir]}"
      exit 2
    end
    script_exporter_cfgs = []
    grafana_agent_scrapes = []

    Dir["#{options[:script_exporter_config_dir]}/**/*.yaml"].each do |f|
      cfg = read_file_to_yaml f
      begin
        validate_config(cfg)
        grafana_agent_scrapes << generate_grafana_agent_scrape(cfg)
        script_exporter_cfgs << generate_partial_script_exporter_config(cfg)
      rescue => error
        @logger.warn "Could not create grafana scrape and script exporter config for #{f} reason: #{error}"
      end
    end

    script_exporter_config = generate_final_script_exporter_config(script_exporter_cfgs)

    # Compare the config on file to the config we just generated, if there is a difference save it and the grafana
    # agent scrapes we generated
    if config_changed(options[:script_exporter_config], script_exporter_config)
      write_config_to_file(options[:script_exporter_config], script_exporter_config)
      run_command('service pharos_script_exporter restart')
    end
    write_scrapes_to_file(options[:grafana_scrape_dir], grafana_agent_scrapes)
  end
end

ScriptExporterSetUp.main if __FILE__ == $PROGRAM_NAME