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

AUTOEncrypt=${3:-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"
INIT_VAULT_URL="./init_vault.sh"
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"
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)
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" "net-tools" "sqlite3")
    
    # Check if lock file exists and delete if needed
    if [ -f /var/lib/dpkg/lock-frontend ]; then
        log "Lock file exists. Deleting lock file..."
        sudo rm /var/lib/dpkg/lock-frontend
    fi
   

    dpkg --configure -a
    if [ $? -ne 0 ]; then
        log "dpkg encountered an error, but we're ignoring it."
    fi 
    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
}



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
        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
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 '/^export MINIO_ACCESSKEY=/d' ~/.bashrc
    echo "export MINIO_ACCESSKEY=$AK" >> ~/.bashrc

    sed -i '/^export 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
function 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
    TRIES=0
    until is_port_open $PORT || [ $TRIES -eq 5 ]; do
        log "Port $PORT is not open. Retrying..."
        sleep 4
        ((TRIES++))
    done

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


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
# Define the cron job command
cron_command="/etc/cron.d/check_docker.sh"

# Check if the cron job already exists
if ! crontab -l | grep -q "$cron_command"; then
    # Add the cron job
    (crontab -l 2>/dev/null; echo "* * * * * $cron_command") | crontab -
    echo "Cron job added successfully."
else
    echo "Cron job already exists."
fi


# 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"

# Check container status and start them
check_and_start_containers() {
    # Check minio container status
    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 start minio >> "$log_file" 2>&1
    fi

    # Check pin_service container status
    if [ ! -f "/opt/facmgr/.pin_service" ]; then
        # .pin_service file doesn't exist, start pin_service container
        pin_service_status=$(docker inspect -f '{{.State.Status}}' pin_service 2>/dev/null)
        if [ "$pin_service_status" != "running" ]; then
            echo "[$(date +"%Y-%m-%d %H:%M:%S")] pin_service container is not running, starting it..." >> "$log_file"
            docker start pin_service >> "$log_file" 2>&1
        fi
    else
        echo "[$(date +"%Y-%m-%d %H:%M:%S")] .pin_service file exists, skipping pin_service container start." >> "$log_file"
    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 start prox_mpd >> "$log_file" 2>&1
    fi
}

# Check and start container status, and log the actions
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
  fi
  if [ "$upgrade" = true ]; then
      create_dir
      restart_docker
      umount_vfs
  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."
