#!/bin/bash
set -e
AK=$1
SK=$2

AUTOEncrypt=${3:-false}
USE_CN_MIRROR=${4:-false}

upgrade=false

# Process the third argument (AUTOEncrypt)
if [ "$3" = "true" ]; then
  AUTOEncrypt=true
fi


for arg in "$@"; do
  if [ "$arg" = "-u" ]; then
    upgrade=true
    break
  fi
done


shift $((OPTIND - 1))

# ... Other script code

if [ "$upgrade" = true ]; then
  echo "Upgrade is set to true."
fi

if [ "$AUTOEncrypt" = true ]; then
  echo "AUTOEncrypt is set to true."
fi
# Define Docker Compose file path and services
COMPOSE_DIR="/opt/facmgr/service"
COMPOSE_FILE="$COMPOSE_DIR/docker-compose.yml"
COMPOSE_URL="./docker-compose.yml"
VAULT_URL="./vault.sh"
INIT_VAULT_CONFIG_URL="./config.hcl"
COMPOSE_FILE_AUTO_ENCRYPT="$COMPOSE_DIR/docker-compose-vault.yml"
SERVICE_FILE_URL="./facmgr_docker.service"
COMPOSE_URL_AUTO_ENCRYPT="./docker-compose-vault.yml"
SERVICE_FILE_PATH="/etc/systemd/system/facmgr_docker.service"
INIT_SCRIPT_URL="./mpd.sql"
VAULT_ENV="/data/vault/config/minio_env.txt"
if [[ ! $AUTOEncrypt == true ]]; then
  SERVICES="prox_mpd pin_service minio minio-token"
else
  #SERVICES="vault kes minio minio-token prox_mpd pin_service"
  SERVICES="vault kes minio minio-token prox_mpd"
fi

LOG_DIR="/var/log/facmgr"
LOG_FILE="$LOG_DIR/docker-compose.log"
DataDIR="/data"
VFS_DIR="$DataDIR/vfs"
MINIO_DIR="$DataDIR/minio"
MINIO_POLICY="/opt/.policy.json"
MINIO_POLICY_URL="./policy.json"
MPD_CONF="$DataDIR/mpd/config"
MDP_DB="$DataDIR/mpd/meta.db"
INIT_SCRIPT="$DataDIR/mpd/mpd.sql"
MPD_LOG="$DataDIR/mpd/log"
PIN_CONFIG="$DataDIR/pin_service/config"
BACKUP_DIR="$DataDIR/mpd/backup"
INIT_VAULT="/data/vault/config"
INIT_DATA="/data/vault"
INIT_VAULT_CONFIG="/data/vault/config/init_vault.sh"
INIT_VAULT_CONFIG_HCL="/data/vault/config/config.hcl"
INIT_MPD_DB_FLAG="/opt/.mpddb"

DIRS=("$COMPOSE_DIR" "$INIT_VAULT" "$MINIO_DIR" "$MPD_CONF" "$MPD_LOG" "$PIN_CONFIG" "$VFS_DIR" "$BACKUP_DIR")
LAST_UPDATE_FILE="/opt/facmgr/last_update.txt"
# Define the port number for each service. Update this according to your services.
declare -A SERVICE_PORTS=(["prox_mpd"]=6007 ["pin_service"]=9091 ["minio"]=9000 ["minio-token"]=9000 ["vault"]=8200 ["kes"]=7373)
MOUNT_POINT=$(df -P "$VFS_DIR" | awk 'END{print $NF}')



create_log() {
    if [ ! -d "$LOG_DIR" ]; then
        echo "Directory $LOG_DIR does not exist. Creating..."
        mkdir -p "$LOG_DIR"
    fi

}

function log {
    echo "$(date +"%Y-%m-%d %H:%M:%S") - $1" | tee -a $LOG_FILE
}

shift $((OPTIND - 1))

if $upgrade; then
  log "Upgrade is set to true."
else
  log "Upgrade is set to false."
fi



check_parm() {
    # Check if AK and SK are set
    if [ -z "$AK" ] || [ -z "$SK" ]; then
        log "AK and/or SK are not set. Please provide them as arguments."
        exit 1
    fi

}

check_root() {
    # Check if script is run as root
    if [[ $EUID -ne 0 ]]; then
        log "This script must be run as root"
        exit 1
    fi
}


package_install() {
    required_tools=("jq" "docker" "docker-compose" "systemctl" "nc" "sqlite3")
    for tool in "${required_tools[@]}"; do
        if dpkg-query -W -f='${Status}' "$tool" 2>/dev/null | grep -q "install ok installed"; then
            log "$tool is already installed. Skipping..."
        else
            log "$tool is missing. Adding to installation list..."
            install_commands+=("apt-get install -y $tool")
        fi
    done

    for cmd in "${install_commands[@]}"; do
        $cmd &
        done
    wait

}

create_dir() {
    # Create necessary directories if they don't exist
    for DIR in "${DIRS[@]}"; do
        if [ ! -d "$DIR" ]; then
            log "Directory $DIR does not exist. Creating..."
            mkdir -p "$DIR"
        fi
    done
}


delete_image() {
    if [[ ! $AUTOEncrypt == true ]]; then
        COMPOSE_FILE="$COMPOSE_FILE"
        COMPOSE_URL="$COMPOSE_URL"
    else
        COMPOSE_FILE="$COMPOSE_FILE_AUTO_ENCRYPT"
        COMPOSE_URL="$COMPOSE_URL_AUTO_ENCRYPT"
    fi

    if [ ! -f "$COMPOSE_FILE" ] || ! cmp -s <(cat "$COMPOSE_URL") "$COMPOSE_FILE"; then
        log "Docker Compose file is missing or outdated. Copying..."
        cp "$COMPOSE_URL" "$COMPOSE_FILE"
    fi


    if [ ! -f "$VAULT_ENV" ]; then
        touch "$VAULT_ENV"
        echo "File $VAULT_ENV created successfully."
    else
        echo "File $VAULT_ENV already exists. Skipping creation."
    fi

    # Check if Docker is running
    if ! systemctl is-active --quiet docker; then
        log "Docker is not running. Starting Docker..."
        systemctl start docker
    fi

    # Check if Docker Compose is running
    if [[ ! $AUTOEncrypt == true ]]; then
        if [ "$(docker-compose -f $COMPOSE_FILE ps -q)" ]; then
            log "Docker Compose is running. Restarting..."
            docker-compose -f $COMPOSE_FILE down
            for service in "${SERVICES[@]}"; do
                docker stop "$service" 2>/dev/null || true
                docker rm "$service" 2>/dev/null || true
                docker rmi "$service" --force 2>/dev/null || true
            done
        fi
    else
        if [ "$(docker-compose -f $COMPOSE_FILE_AUTO_ENCRYPT ps -q)" ]; then
            log "Docker Compose is running. Restarting..."
            docker-compose -f $COMPOSE_FILE_AUTO_ENCRYPT down
            for service in "${SERVICES[@]}"; do
                docker stop "$service" 2>/dev/null || true
                docker rm "$service" 2>/dev/null || true
                docker rmi "$service" --force 2>/dev/null || true
            done
        fi
    fi


    for service in "${SERVICES[@]}"; do
        if docker ps -aq --filter "name=$service" >/dev/null; then
            docker stop "$service" 2>/dev/null || true
            docker rm "$service" 2>/dev/null || true
            docker rmi "$service" --force 2>/dev/null || true
            echo "Container $service stopped, removed, and image deleted."
        else
            echo "Container $service does not exist."
        fi
    done


    for image in "${SERVICES[@]}"; do
        if docker images -q "$image" >/dev/null; then
            docker rmi "$image" --force 2>/dev/null || true
            echo "Image $image deleted."
        else
            echo "Image $image does not exist."
        fi
    done
}

umount_vfs(){
# Unmount directory if necessary
    if [ "$VFS_DIR" = "$MOUNT_POINT" ]; then
        log "Directory $VFS_DIR is a mount point. Unmounting..."
        umount -l "$VFS_DIR" 2>/dev/null || true
    fi
}


init_vault(){ 
    cp -r $MINIO_POLICY_URL $MINIO_POLICY

    if [[ $AUTOEncrypt == true ]]; then
        sed -i "s,{{ use_cn_mirror }},$USE_CN_MIRROR,g" $VAULT_URL
        cp -r $VAULT_URL $COMPOSE_DIR/vault.sh
        source $COMPOSE_DIR/vault.sh
        cp -r $INIT_VAULT_CONFIG_URL $INIT_VAULT_CONFIG_HCL
        chmod +x $INIT_VAULT_CONFIG
    fi
}


pull_image() {
      	
    # Get the last update time
    if [ -f $LAST_UPDATE_FILE ]; then
        LAST_UPDATE=$(cat $LAST_UPDATE_FILE)
    else
        LAST_UPDATE=0
    fi
    export DOCKER_CLIENT_TIMEOUT=600
    export COMPOSE_HTTP_TIMEOUT=600
    log "More than a week has passed since the last update. Checking for updates..."
    systemctl restart docker
    docker builder prune -f
    log "docker builder prune successful"

    if [[ ! $AUTOEncrypt == true ]]; then
        COMPOSE_HTTP_TIMEOUT=600 docker-compose -f $COMPOSE_FILE pull
        date +%s > $LAST_UPDATE_FILE
    else
        touch $VAULT_ENV
        COMPOSE_HTTP_TIMEOUT=600 docker-compose -f $COMPOSE_FILE_AUTO_ENCRYPT pull
        date +%s > $LAST_UPDATE_FILE
    fi

    log "docker-compose pull image"

    if [ $? -eq 124 ]; then
        log "Pull command timed out. Performing individual image pulls..."
        exit 1
        # Perform individual image pulls using docker pull
        # Replace the <service-name> with your service name
    fi
}

init_mpd_db(){

    if [ ! -f "$INIT_MPD_DB_FLAG" ]; then
        log "Database not found, initializing..."
        cp -r $INIT_SCRIPT_URL $INIT_SCRIPT
        sqlite3 "$MDP_DB" < "$INIT_SCRIPT"
        touch $INIT_MPD_DB_FLAG
    fi
}

# Check if a port is open
#is_port_open() {
#    if [[ $1 -gt 0 ]]; then
#        #nc -z localhost "$1" >/dev/null 2>&1
#        nc -z localhost "$1"
#    fi
#}


function is_port_open {
    local port=$1
    if [[ $port -gt 0 ]]; then
            netstat -anp | grep ":$port\s" | grep "LISTEN"
    fi
}

# Check if a service is healthy
function is_service_healthy {
    docker inspect --format "{{json .State.Health.Status }}" $1
    docker inspect --format "{{json .State.Health.Status }}" $1 | grep -q "healthy"
}


env_set() {
    if [ -n "$MINIO_ACCESSKEY" ]; then
        unset MINIO_ACCESSKEY
    fi

    if [ -n "$MINIO_SECRETKEY" ]; then
        unset MINIO_SECRETKEY
    fi

    export MINIO_ACCESSKEY=$AK
    export MINIO_SECRETKEY=$SK

    sed -i '/^MINIO_ACCESSKEY=/d' ~/.bashrc
    echo "export MINIO_ACCESSKEY=$AK" >> ~/.bashrc

    sed -i '/^MINIO_SECRETKEY=/d' ~/.bashrc
    echo "export MINIO_SECRETKEY=$SK" >> ~/.bashrc

    if [ -f ~/.bashrc ]; then
        source ~/.bashrc
    fi

    #env_file="$(pwd)/.env"
}

minio_env(){

  env_file="/opt/facmgr_install/.env"
  if [ -f "$env_file" ] && grep -q "MINIO_ACCESSKEY" "$env_file" && grep -q "MINIO_SECRETKEY" "$env_file"; then
    echo "Environment variables already exist in $env_file. Skipping writing."
  else
    echo "MINIO_ACCESSKEY=$AK" > "$env_file"
    echo "MINIO_SECRETKEY=$SK" >> "$env_file"
    echo "Environment variables have been written to $env_file."
  fi
}

# Start a service and check its status

start_service() {
    SERVICE_NAME=$1
    PORT=${SERVICE_PORTS[$SERVICE_NAME]}

    log "Starting $SERVICE_NAME..."
    docker-compose -f $COMPOSE_FILE up -d $SERVICE_NAME

    log "Waiting for $SERVICE_NAME to start..."
    sleep 10

    # Check port status with retries
    MAX_RETRIES=30
    TRIES=0
    RETRY_FAILED_COUNT=0
    RETRY_INTERVAL=10
    while [ $TRIES -lt $MAX_RETRIES ]; do
        if is_port_open $PORT; then
            log "$SERVICE_NAME started and is healthy."
            break
        else
            echo "Port $PORT is not open. Retrying..."
            sleep $RETRY_INTERVAL
	    TRIES=`expr $TRIES + 1`;
	    RETRY_FAILED_COUNT=`expr $RETRY_FAILED_COUNT + 1`;
        fi
    done

    if [ $RETRY_FAILED_COUNT -eq $MAX_RETRIES ]; then
        log "Port $PORT is still not open after $TRIES attempts. Check the log file for details."
        exit 1
    fi
}







run_service() {
    # Start services
    for SERVICE in $SERVICES; do
        start_service $SERVICE
    done

    # Check if the systemctl service exists
    if ! systemctl list-units --full --all | grep -Fq "facmgr_docker.service"; then
        log "facmgr_docker.service does not exist. Downloading..."
        cp -r  $SERVICE_FILE_URL $SERV $SERVICE_FILE_PATH
        systemctl daemon-reload
        systemctl enable facmgr_docker.service
    fi

    log "All services started successfully."
}

cron_job() {
    # Define cron jobs
    BACKUP_JOB_IDENTIFIER="# FACMGR_DB_BACKUP_JOB"
    CRON_JOB_BACKUP="$BACKUP_JOB_IDENTIFIER\n0 1 * * * /usr/bin/env sqlite3 $MDP_DB '.backup '$BACKUP_DIR/db_backup_'\$(date +\%Y\%m\%d\%H\%M)'.sqlite'"

    CLEANUP_JOB_IDENTIFIER="# FACMGR_DB_CLEANUP_JOB"
    CRON_JOB_CLEANUP="$CLEANUP_JOB_IDENTIFIER\n0 2 * * * /usr/bin/env find $BACKUP_DIR -name 'db_backup_*.sqlite' -mtime +30 -delete"

    # Add the backup cron job if it does not exist
    if ! crontab -l | grep -q "$BACKUP_JOB_IDENTIFIER"; then
        (crontab -l 2>/dev/null; echo -e "$CRON_JOB_BACKUP") | crontab -
    fi

    # Add the cleanup cron job if it does not exist
    if ! crontab -l | grep -q "$CLEANUP_JOB_IDENTIFIER"; then
        (crontab -l 2>/dev/null; echo -e "$CRON_JOB_CLEANUP") | crontab -
    fi
}


function add_cyfs_toml {

if [ ! -d "/cyfs/etc/app-manager" ]; then
    mkdir -p "/cyfs/etc/app-manager"
fi

cat > "/cyfs/etc/app-manager/app-manager.toml" <<EOF
[config]
sandbox = "default"

[app]
source = "user"
EOF


}

cron_start_docker(){

# Add the cron job

(crontab -l 2>/dev/null | grep -v "/etc/cron.d/check_docker.sh" ; echo "* * * * * /etc/cron.d/check_docker.sh") | crontab -

# Create the script file and write the content
script_file="/etc/cron.d/check_docker.sh"
cat << 'EOF' > "$script_file"
#!/bin/bash

log_file="/var/log/start_docker.log"

declare -A SERVICE_PORTS=(["prox_mpd"]=6007 ["pin_service"]=9091 ["minio"]=9000 ["minio-token"]=9000 ["vault"]=8200 ["kes"]=7373)
COMPOSE_DIR="/opt/facmgr_install/"
COMPOSE_FILE="$COMPOSE_DIR/docker-compose-vault.yml"


function is_port_open {
    local port=$1
    if [[ $port -gt 0 ]]; then
            netstat -anp | grep ":$port\s" | grep "LISTEN"
    fi
}


start_service() {
    SERVICE_NAME=$1
    PORT=${SERVICE_PORTS[$SERVICE_NAME]}

    echo "Starting $SERVICE_NAME..." >> "$log_file"
    cd $COMPOSE_DIR
    docker-compose -f $COMPOSE_FILE up -d $SERVICE_NAME

    echo "Waiting for $SERVICE_NAME to start..." >> "$log_file"

    sleep 10

    # Check port status with retries
    MAX_RETRIES=30
    TRIES=0
    RETRY_FAILED_COUNT=0
    RETRY_INTERVAL=10
    while [ $TRIES -lt $MAX_RETRIES ]; do
        if is_port_open $PORT; then
            echo "$SERVICE_NAME started and is healthy." >> "$log_file"
            break
        else
            echo "Port $PORT is not open. Retrying..." >> "$log_file"
            sleep $RETRY_INTERVAL
            TRIES=`expr $TRIES + 1`;
            RETRY_FAILED_COUNT=`expr $RETRY_FAILED_COUNT + 1`;
        fi
    done

    if [ $RETRY_FAILED_COUNT -eq $MAX_RETRIES ]; then
        echo "Port $PORT is still not open after $TRIES attempts. Check the log file for details." >> "$log_file"
        exit 1
    fi
}



# Check container status and start them
check_and_start_containers() {
    # Check minio container status

    vault_status=$(docker inspect -f '{{.State.Status}}' vault 2>/dev/null)
    if [ "$vault_status" != "running" ]; then
        echo "[$(date +"%Y-%m-%d %H:%M:%S")] vault container is not running, starting it..." >> "$log_file"
        docker rm vault >> "$log_file" 2>&1
	docker stop vault >> "$log_file" 2>&1
        start_service vault >> "$log_file" 2>&1
    fi


    kes_status=$(docker inspect -f '{{.State.Status}}' kes 2>/dev/null)
    if [ "$kes_status" != "running" ]; then
        echo "[$(date +"%Y-%m-%d %H:%M:%S")] kes container is not running, starting it..." >> "$log_file"
        rm -rf /data/kes/certs/*
        rm -rf /data/minio/certs/* 
	docker rm kes >> "$log_file" 2>&1
        docker stop kes >> "$log_file" 2>&1
        start_service kes >> "$log_file" 2>&1
    fi

    minio_status=$(docker inspect -f '{{.State.Status}}' minio 2>/dev/null)
    if [ "$minio_status" != "running" ]; then
        echo "[$(date +"%Y-%m-%d %H:%M:%S")] minio container is not running, starting it..." >> "$log_file"
        docker rm minio >> "$log_file" 2>&1
        docker stop minio >> "$log_file" 2>&1
        start_service minio >> "$log_file" 2>&1
    fi


    # Check prox_mpd container status
    prox_mpd_status=$(docker inspect -f '{{.State.Status}}' prox_mpd 2>/dev/null)
    if [ "$prox_mpd_status" != "running" ]; then
        echo "[$(date +"%Y-%m-%d %H:%M:%S")] prox_mpd container is not running, starting it..." >> "$log_file"
        umount /data/vfs -l
	docker rm prox_mpd >> "$log_file" 2>&1
        docker stop prox_mpd >> "$log_file" 2>&1
        start_service prox_mpd >> "$log_file" 2>&1
    fi
}


check_and_start_containers


echo "[$(date +"%Y-%m-%d %H:%M:%S")] Container check and start script executed" >> "$log_file"
EOF

# Make the script executable
chmod +x "$script_file"
systemctl restart cron.service
log "Cron job has been added, and the script file $script_file has been created."

}

restart_docker() {
    systemctl restart docker
    log "Start to restart Docker"
}


main() {
  create_log
  check_parm
  check_root
  if [ "$upgrade" = false ]; then
      rm -rf $INIT_DATA
      rm -rf $INIT_MPD_DB_FLAG
      package_install
      umount_vfs
      create_dir
      env_set
      src_dir="/opt/facmgr_install"
      dst_mpd_dir="/data/mpd/config"
      dst_pin_service_dir="/data/pin_service"

      if [ ! -f "$dst_mpd_dir/config.json" ] && [ -f "$src_dir/config.json" ]; then
        cp "$src_dir/config.json" "$dst_mpd_dir/"
        echo "Copied config.json to $dst_mpd_dir"
      fi

      if [ ! -f "$dst_pin_service_dir/config.toml" ] && [ -f "$src_dir/config.toml" ]; then
        cp "$src_dir/config.toml" "$dst_pin_service_dir/"
        echo "Copied config.toml to $dst_pin_service_dir"
      fi
  fi
  if [ "$upgrade" = true ]; then
      create_dir
      restart_docker
      umount_vfs
      [ -f /etc/cron.d/check_docker.sh ] && rm /etc/cron.d/check_docker.sh  && systemctl restart cron

  fi

 
  minio_env
  init_vault
  delete_image
  pull_image
  run_service
  init_mpd_db
  cron_job
  add_cyfs_toml
  cron_start_docker
}
main


log "Cron jobs for database backup and cleanup are created."
