#!/bin/bash
set -e

ENABLE_KNOWLEDGE_ENV=${2:-true}
OPENAI_API_KEY_ENV=${1:-sk-123}
API_TOKEN_ENV=${3:-123}

upgrade=false

if [ "$2" = "false" ]; then
    ENABLE_KNOWLEDGE_ENV=false
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



# Define Docker Compose file path and services
COMPOSE_DIR="/opt/prox_ai_install"
COMPOSE_FILE="$COMPOSE_DIR/docker-compose.yml"
SERVICE_FILE_URL="./prox_ai_docker.service"
SERVICE_FILE_PATH="/etc/systemd/system/prox_ai_docker.service"
SERVICES="prox_ai_lite"

LOG_DIR="/var/log/facmgr"
LOG_FILE="$LOG_DIR/prox_ai_lite.log"

LAST_UPDATE_FILE="/opt/facmgr/last_prox_ai_update.txt"
# Define the port number for each service. Update this according to your services.
declare -A SERVICE_PORTS=(["prox_ai_lite"]=9009)


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

    if [ ! -f "$LOG_FILE" ]; then
        touch "$LOG_FILE"
    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_root() {
    # Check if script is run as root
    if [[ $EUID -ne 0 ]]; then
        log "This script must be run as root"
        exit 1
    fi
}



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() {
    COMPOSE_FILE="$COMPOSE_FILE"

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

    if [ "$(docker-compose -f $COMPOSE_FILE ps -q)" ]; then
        log "Docker Compose is running. Restarting..."
        docker-compose -f $COMPOSE_FILE down
    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
            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
}

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"

    COMPOSE_HTTP_TIMEOUT=600 docker-compose -f $COMPOSE_FILE pull
    date +%s > $LAST_UPDATE_FILE

    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
}

function is_port_open {
    if [[ $1 -gt 0 ]]; then
        nc -z localhost "$1" >/dev/null 2>&1
    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 "$OPENAI_API_KEY" ]; then
        unset OPENAI_API_KEY
    fi

    if [ -n "$ENABLE_KNOWLEDGE" ]; then
        unset ENABLE_KNOWLEDGE
    fi
    if [ -n "$API_TOKEN" ]; then
        unset API_TOKEN
    fi
    export OPENAI_API_KEY=$OPENAI_API_KEY_ENV
    export ENABLE_KNOWLEDGE=$ENABLE_KNOWLEDGE_ENV
    export API_TOKEN=$API_TOKEN_ENV

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

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

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

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

}


prox_env(){

  env_file="/opt/prox_ai_install/.env"
  if [ -f "$env_file" ] && grep -q "OPENAI_API_KEY" "$env_file" && grep -q "ENABLE_KNOWLEDGE" "$env_file"; then
    echo "Environment variables already exist in $env_file. Skipping writing."
  else
    echo "OPENAI_API_KEY=$OPENAI_API_KEY_ENV" > "$env_file"
    echo "ENABLE_KNOWLEDGE=$ENABLE_KNOWLEDGE_ENV" >> "$env_file"
    echo "API_TOKEN=$API_TOKEN_ENV" >> "$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 "prox_ai_docker.service"; then
        log "prox_ai_docker.service does not exist. Downloading..."
        cp -r  $SERVICE_FILE_URL $SERV $SERVICE_FILE_PATH
        systemctl daemon-reload
        systemctl enable prox_ai_docker.service
    fi

    log "All services started successfully."
}


cron_start_docker(){



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



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

log_file="/var/log/start_prox_ai_lite_docker.log"

# Check container status and start them
check_and_start_containers() {

        prox_ai_lite_service_status=$(docker inspect -f '{{.State.Status}}' prox_ai_lite 2>/dev/null)
        if [ "$prox_ai_lite_status" != "running" ]; then
            echo "[$(date +"%Y-%m-%d %H:%M:%S")] prox_ai_lite container is not running, starting it..." >> "$log_file"
            docker start prox_ai_lite >> "$log_file" 2>&1
        fi
    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_root
  if [ "$upgrade" = false ]; then
      env_set
  fi
  #if [ "$upgrade" = true ]; then
  #    restart_docker
  #fi

  prox_env
  delete_image
  pull_image
  run_service
  cron_start_docker
}
main

log "Cron jobs for prox ai are created."
