#! /bin/bash

echo "
  ____             __          _       _                
 / ___|    __ _   / _|   ___  | |     (_)  _ __     ___ 
 \___ \   / _\` | | |_   / _ \ | |     | | | '_ \   / _ \\
  ___) | | (_| | |  _| |  __/ | |___  | | | | | | |  __/
 |____/   \__,_| |_|    \___| |_____| |_| |_| |_|  \___|
"

export STREAM=${STREAM:-0}
export CDN=${CDN:-1}

RELEASE="${RELEASE}"
if [ -n "$RELEASE" ]; then
    RELEASE="-$RELEASE"
fi

qrcode() {
    echo "https://discord.gg/wyshSVuvxC"
    echo
    echo "Visit the above address to join the SafeLine discord group"
}

command_exists() {
    command -v "$1" 2>&1
}

check_container_health() {
    local container_name=$1
    local max_retry=30
    local retry=0
    local health_status="unhealthy"
    echo "Waiting for $container_name to be healthy"
    while [[ "$health_status" == "unhealthy" && $retry -lt $max_retry ]]; do
        health_status=$(docker inspect --format='{{.State.Health.Status}}' $container_name 2>/dev/null || echo 'unhealthy')
        sleep 5
        retry=$((retry+1))
    done
    if [[ "$health_status" == "unhealthy" ]]; then
        abort "Container $container_name is unhealthy"
    fi
    echo "Container $container_name is healthy"
}

space_left() {
    dir="$1"
    while [ ! -d "$dir" ]; do
        dir=$(dirname "$dir")
    done
    echo $(df -h "$dir" --output='avail' | tail -n 1)
}

confirm() {
    echo -e -n "\033[34m[SafeLine] $* \033[1;36m(Y/n)\033[0m"
    read -n 1 -s opt

    [[ "$opt" == $'\n' ]] || echo

    case "$opt" in
    'y' | 'Y') return 0 ;;
    'n' | 'N') return 1 ;;
    *) confirm "$1" ;;
    esac
}

info() {
    echo -e "\033[37m[SafeLine] $*\033[0m"
}

warning() {
    echo -e "\033[33m[SafeLine] $*\033[0m"
}

abort() {
    qrcode
    echo -e "\033[31m[SafeLine] $*\033[0m"
    exit 1
}

trap 'onexit' INT
onexit() {
    echo
    abort "User terminates the installation process"
}

if [ -z "$BASH" ]; then
    abort "Please use bash to execute the current script, please refer to the latest technical documentation https://waf.chaitin.com/"
fi

if [ ! -t 0 ]; then
    abort "STDIN is not a standard input device, please refer to the latest technical documentation https://waf.chaitin.com/"
fi

if [ "$#" -ne "0" ]; then
    abort "The current script does not require any parameters, please refer to the latest technical documentation https://waf.chaitin.com/"
fi

if [ "$EUID" -ne "0" ]; then
    abort "Please run with root privileges"
fi
info "The script calling method is confirmed to be normal"

if [ -z `command_exists docker` ]; then
    warning "Docker is not installed"
    if confirm "Do you need to install Docker automatically?"; then
        curl -sSLk https://get.docker.com/ | bash
        if [ $? -ne "0" ]; then
            abort "Docker installation failed"
        fi
        info "Docker installation finished"
    else
        abort "Installation has been terminate"
    fi
fi
info "Docker is installed at: '`command -v docker`'"

start_docker
docker version > /dev/null 2>&1
if [ $? -ne "0" ]; then
    abort "Docker service is inactive"
fi
info "Docker service is active"

compose_command="docker compose"
if $compose_command version; then
    info "Docker Compose Plugin is found"
else
    warning "Docker Compose Plugin is not found"
    compose_command="docker-compose"
    if [ -z `command_exists "docker-compose"` ]; then
        warning "docker-compose is not found"
        if confirm "Do you need to install Docker Compose Plugin automatically"; then
            curl -sSLk https://get.docker.com/ | bash
            if [ $? -ne "0" ]; then
                abort "Docker Compose Plugin installation failed"
            fi
            info "Docker Compose Plugin installation finished"
            compose_command="docker compose"
        else
            abort "Installation has been terminate"
        fi
    else
        info "docker-compose is installed at: '`command -v docker-compose`'"
    fi
fi


container_id=$(docker ps -n 1 --filter name=.*safeline-mgt.* --format '{{.ID}}')
safeline_path=$(docker inspect --format '{{index .Config.Labels "com.docker.compose.project.working_dir"}}' $container_id)

while [ -z "$safeline_path" ]; do
    echo -e -n "\033[34m[SafeLine] No running SafeLine was found, input the SafeLine directory. (blank means '/data/safeline'): \033[0m"
    read input_path
    [[ -z "$input_path" ]] && input_path='/data/safeline'

    if [[ ! $input_path == /* ]]; then
        warning "'$input_path' is not valid path"
        continue
    fi

    safeline_path=$input_path
done

cd "$safeline_path"

grep COLLIE .env >/dev/null 2>&1
if [ $? -eq "0" ]; then
    abort "It is detected that your environment is installed through the Muyun Host Assistant. Please use the Muyun Host Assistant-Application Market to upgrade"
fi

compose_name=$(ls docker-compose.yaml compose.yaml 2>/dev/null)
compose_path=$safeline_path/$compose_name

mv $compose_name $compose_name.old || true

curl "https://waf.chaitin.com/release/latest/compose.yaml" -sSLk -o $compose_name
curl "https://waf.chaitin.com/release/latest/reset_tengine.sh" -sSLk -o reset_tengine.sh

if [ $? -ne "0" ]; then
    abort "Download compose.yaml failed"
fi
info "Download compose.yaml successful"

if [ $STREAM -eq 1 ]; then
    sed -i "s/IMAGE_TAG=.*/IMAGE_TAG=latest-stream/g" ".env"
else
    sed -i "s/IMAGE_TAG=.*/IMAGE_TAG=latest/g" ".env"
fi

grep "SAFELINE_DIR" ".env" >/dev/null || echo "SAFELINE_DIR=$(pwd)" >>".env"

if [ $STREAM -eq 1 ]; then
    grep "IMAGE_TAG" ".env" >/dev/null || echo "IMAGE_TAG=latest-stream" >>".env"
else
    grep "IMAGE_TAG" ".env" >/dev/null || echo "IMAGE_TAG=latest" >>".env"
fi

grep "MGT_PORT" ".env" >/dev/null || echo "MGT_PORT=9443" >>".env"
grep "POSTGRES_PASSWORD" ".env" >/dev/null || echo "POSTGRES_PASSWORD=$(LC_ALL=C tr -dc A-Za-z0-9 </dev/urandom | head -c 32)" >>".env"
grep "SUBNET_PREFIX" ".env" >/dev/null || echo "SUBNET_PREFIX=172.22.222" >>".env"

ARCH=$(uname -m)
if [[ "$ARCH" =~ "aarch" || "$ARCH" =~ "arm" ]]; then
    grep "ARCH_SUFFIX" ".env" >/dev/null || echo -e "\nARCH_SUFFIX=-arm" >>".env"
else
    grep "ARCH_SUFFIX" ".env" >/dev/null || echo -e "\nARCH_SUFFIX=" >>".env"
fi

grep "RELEASE" ".env" >/dev/null || echo -e "\nRELEASE=$RELEASE" >>".env"
grep "REGION" ".env" >/dev/null || echo -e "\nREGION=-g" >>".env"

info "Downloading Docker image"

$compose_command pull
if [ $? -ne "0" ]; then
    abort "Docker container starts failed"
fi


docker rm -f safeline-redis &>/dev/null

$compose_command down --remove-orphans && $compose_command up -d
if [ $? -ne "0" ]; then
    abort "Docker container starts failed"
fi

qrcode

check_container_health safeline-pg
check_container_health safeline-mgt
docker exec safeline-mgt /app/mgt-cli reset-admin --once

warning "SafeLine WAF has been installed, visit the following address to Web admin console"
warning "https://0.0.0.0:9443/"

