#!/bin/bash

WORK_DIR=/opt/nezha/nz-agent

fn_init_nz_agent_service(){
  cat > /usr/lib/systemd/system/nz-agent.service <<EOF
[Unit]
Description=nz-agent
After=network.target
[Service]
WorkingDirectory=/opt/nezha/nz-agent
ExecStart=/opt/nezha/nz-agent/xjar ${JAVA_EXE} -Dnz-agent.dir=/opt/nezha/nz-agent -jar /opt/nezha/nz-agent/nz-agent.xjar
RestartSec=5s
Restart=always
[Install]
WantedBy=multi-user.target
EOF
}

fn_init_prometheus_service(){
    cat > /usr/lib/systemd/system/prometheus.service <<"EOF"
[Unit]
Description=prometheus
After=network.target
[Service]
WorkingDirectory=/opt/nezha/nz-agent/prometheus
EnvironmentFile=/opt/nezha/nz-agent/prometheus/config.conf
ExecStart=/opt/nezha/nz-agent/prometheus/prometheus $OPTION
RestartSec=10s
Restart=always
LimitNOFILE=524288
[Install]
WantedBy=multi-user.target
EOF
}

fn_init_snmp_exporter_service(){
    cat > /usr/lib/systemd/system/snmp-exporter.service <<"EOF"
[Unit]
Description=snmp-exporter
After=network.target
[Service]
EnvironmentFile=/opt/nezha/nz-agent/snmp_exporter/config.conf
ExecStart=/opt/nezha/nz-agent/snmp_exporter/snmp_exporter $OPTION
RestartSec=10s
Restart=always
LimitNOFILE=524288
[Install]
WantedBy=multi-user.target
EOF
}

fn_init_blackbox_exporter_service(){
    cat > /usr/lib/systemd/system/blackbox-exporter.service <<"EOF"
[Unit]
Description=blackbox_exporter
After=network.target
[Service]
WorkingDirectory=/opt/nezha/nz-agent/blackbox_exporter
EnvironmentFile=/opt/nezha/nz-agent/blackbox_exporter/config.conf
ExecStart=/opt/nezha/nz-agent/blackbox_exporter/blackbox_exporter $OPTION
RestartSec=10s
Restart=always
LimitNOFILE=524288
[Install]
WantedBy=multi-user.target
EOF
}

fn_init_cortex_service(){
  cat > /usr/lib/systemd/system/cortex.service <<"EOF"
[Unit]
Description=cortex
After=network.target
[Service]
WorkingDirectory=/opt/nezha/nz-agent/cortex
EnvironmentFile=/opt/nezha/nz-agent/cortex/config.conf
ExecStart=/opt/nezha/nz-agent/cortex/cortex-linux-amd64 $OPTION
RestartSec=10s
Restart=always
LimitNOFILE=524288
[Install]
WantedBy=multi-user.target
EOF
}

fn_init_loki_service(){
  cat > /usr/lib/systemd/system/loki.service <<"EOF"
[Unit]
Description=loki
After=network.target
[Service]
WorkingDirectory=/opt/nezha/nz-agent/loki
EnvironmentFile=/opt/nezha/nz-agent/loki/config.conf
ExecStart=/opt/nezha/nz-agent/loki/loki-linux-amd64 $OPTION
RestartSec=10s
Restart=always
LimitNOFILE=524288
[Install]
WantedBy=multi-user.target
EOF
}

fn_init_config(){
  # nz-agent component
  chmod +x /opt/nezha/nz-agent/xjar
  
  # prometheus component
  cat > /opt/nezha/nz-agent/prometheus/config.conf <<EOF
OPTION=" --config.file='/opt/nezha/nz-agent/prometheus/prometheus.yml' --web.listen-address='127.0.0.1:10091' --storage.tsdb.path='/opt/nezha/data' --web.enable-admin-api --web.enable-lifecycle --storage.tsdb.retention.time=15d "
EOF

  # snmp-exporter component
  cat >/opt/nezha/nz-agent/snmp_exporter/config.conf <<EOF
OPTION="--web.listen-address='127.0.0.1:19116' --config.file='/opt/nezha/nz-agent/snmp_exporter/snmp.yml' "
EOF

  # blackbox-exporter component
  cat >/opt/nezha/nz-agent/blackbox_exporter/config.conf <<EOF
OPTION="--web.listen-address='127.0.0.1:19115' --config.file='/opt/nezha/nz-agent/blackbox_exporter/blackbox.yml' "
EOF

  # cortex component
  chmod +x /opt/nezha/nz-agent/cortex/cortex-linux-amd64
  cat >/opt/nezha/nz-agent/cortex/config.conf <<EOF
OPTION=" -config.file=/opt/nezha/nz-agent/cortex/cortex.yml"
EOF

  # loki component
  chmod +x /opt/nezha/nz-agent/loki/loki-linux-amd64
  cat >/opt/nezha/nz-agent/loki/config.conf <<EOF
OPTION=" -config.file=/opt/nezha/nz-agent/loki/loki.yml"
EOF
}

compareMD5(){
  if [ ! -f $1 ] || [ ! -f $2 ];then
    echo 1
    return 1
  fi

  local f1MD5=`md5sum $1|awk '{print $1}'`
  local f2MD5=`md5sum $2|awk '{print $1}'`

  if [ ${f1MD5} = ${f2MD5} ];then
      echo 0
      return 0
  else
      echo 1
      return 1
  fi
}

# 还原配置文件
# $1 需要还原的文件目录 /opt/nezha/nz-agent/config  /opt/nezha/nz-agent/prometheus ...
# $2 备份文件目录       /tmp/nezha/nz-agent/config /tmp/nezha/nz-agent/prometheus ...
restoreComponentConfig(){
  if [ ! -d $1 ] || [ ! -d $2 ];then
    echo "Directory does not exist"
    return 1
  fi

  for i in $(ls $2);do
    if [ 1 -eq `compareMD5 $1/$i $2/$i` ];then
      echo 'return config file '$1/${i}
      cp -f $2/$i $1/$i
    fi
  done
}

# Find Java
if [[ -n "$WORK_DIR/jdk" ]] && [[ -x "$WORK_DIR/jdk/bin/java" ]]; then
    JAVA_EXE="$WORK_DIR/jdk/bin/java"
elif [[ -n "$JAVA_HOME" ]] && [[ -x "$JAVA_HOME/bin/java" ]]; then
    JAVA_EXE="$JAVA_HOME/bin/java"
elif type -p java > /dev/null 2>&1; then
    JAVA_EXE=$(type -p java)
elif [[ -x "/usr/bin/java" ]];  then
    JAVA_EXE="/usr/bin/java"
else
    echo "Unable to find Java"
    JAVA_EXE="java"
#    exit 1
fi

# install
if [ 1 -eq $1 ];then
  fn_init_config
  fn_init_nz_agent_service
  fn_init_prometheus_service
  fn_init_snmp_exporter_service
  fn_init_blackbox_exporter_service
  fn_init_cortex_service
  fn_init_loki_service
 
  systemctl daemon-reload
  systemctl enable nz-agent && systemctl restart nz-agent
  systemctl enable prometheus && systemctl restart prometheus
  systemctl enable snmp-exporter && systemctl restart snmp-exporter
  systemctl enable blackbox-exporter && systemctl restart blackbox-exporter
  systemctl enable cortex && systemctl restart cortex
  systemctl enable loki && systemctl restart loki
  echo 'install nz-agent success !'
fi

# update
if [ 2 -eq $1 ];then
  fn_init_config
  # nz-agent component 
  if [ ! -f "/usr/lib/systemd/system/nz-agent.service" ];then
  	fn_init_nz_agent_service
  fi
  
  # prometheus component 
  if [ ! -f "/usr/lib/systemd/system/prometheus.service" ];then
	fn_init_prometheus_service
  fi

  # snmp-exporter component 
  if [ ! -f "/usr/lib/systemd/system/snmp-exporter.service" ];then
	fn_init_snmp_exporter_service
  fi

  # blackbox-exporter component 
  if [ ! -f "/usr/lib/systemd/system/blackbox-exporter.service" ];then
	fn_init_blackbox_exporter_service
  fi

  # cortex component 
  if [ ! -f "/usr/lib/systemd/system/cortex.service" ];then
	fn_init_cortex_service
  fi

  # cortex component
  if [ ! -f "/usr/lib/systemd/system/loki.service" ];then
	fn_init_loki_service
  fi
  
  # restore nz-agent config
  AGENT_PATH=/opt/nezha/nz-agent
  TMP_PATH=/tmp/nezha/nz-agent
  echo 'move backup config file...'
  if [ -d $TMP_PATH ];then
	#for i in $(ls $AGENT_PATH/config);do
	#	if [ 1 -eq `compareMD5 $AGENT_PATH/config/$i $TMP_PATH/config/$i` ];then
	#		echo 'return config file '${i}
	#		#mv -f $AGENT_PATH/config/$i $AGENT_PATH/config/$i
	#		cp -f $TMP_PATH/config/$i $AGENT_PATH/config/$i
	#		#cat $AGENT_PATH/config/token.auth'.oldconfig' > $AGENT_PATH/config/token.auth && rm -f $AGENT_PATH/config/token.auth'.oldconfig'
	#	fi
	#done

	# nz-agent config
	restoreComponentConfig $AGENT_PATH/config $TMP_PATH/config
    # blackbox_exporter
    restoreComponentConfig $AGENT_PATH/blackbox_exporter $TMP_PATH/blackbox_exporter
    # snmp_exporter
    restoreComponentConfig $AGENT_PATH/snmp_exporter $TMP_PATH/snmp_exporter
    # cortex
    restoreComponentConfig $AGENT_PATH/cortex $TMP_PATH/cortex
    # loki
    restoreComponentConfig $AGENT_PATH/loki $TMP_PATH/loki
    # prometheus
    restoreComponentConfig $AGENT_PATH/prometheus $TMP_PATH/prometheus
  fi

  # restart all components
  systemctl daemon-reload
  systemctl enable nz-agent && systemctl restart nz-agent
  systemctl enable prometheus && systemctl restart prometheus
  systemctl enable snmp-exporter && systemctl restart snmp-exporter
  systemctl enable blackbox-exporter && systemctl restart blackbox-exporter
  systemctl enable cortex && systemctl restart cortex
  systemctl enable loki && systemctl restart loki
  echo 'update nz-agent success !'
fi