#!/bin/bash

#################################################################################################
###  
#################################################################################################
VERSION="V3.0.5"

# Retry configuration
MAX_RETRIES=3
RETRY_DELAY=5

declare -A IMAGES
# Define images and their versions to check
# Target images and tags
# Define images and their versions
IMAGES["luscis/openlan"]="V3.0.5"
# IMAGES["prom/node-exporter"]="v1.6.1"  # 注释掉 node-exporter

# Target images and tags
IRISNODE_IMAGE="luscis/irisnode"
IRISNODE_TAG="V3.0.5" 
IRISNODE_VERSION="V3.0.5"
IRISNODE_IMAGE_ID="66642367cdb2"

# NODE_EXPORTER_IMAGE="prom/node-exporter"  # 注释掉
# NODE_EXPORTER_TAG="v1.6.1"  # 注释掉
# NODE_EXPORTER_VERSION="v1.6.1"  # 注释掉
# NODE_EXPORTER_IMAGE_ID="458e026e6aa6"  # 注释掉

##########################################################################################
# Function definitions

# Generate 12-character random secret with numbers and lowercase letters
generate_random_secret() {
    local charset="0123456789abcdefghijklmnopqrstuvwxyz"
    local result=""
    for ((i = 0; i < 12; i++)); do
        index=$((RANDOM % ${#charset}))
        result="${result}${charset:$index:1}"
    done
    echo "$result"
}

# Command execution with retry and timeout handling
execute_with_retry() {
    local cmd="$1"
    local description="$2"
    local timeout="${3:-300}"  # Default timeout 5 minutes
    
    echo "Starting: $description"
    
    for ((attempt=1; attempt<=MAX_RETRIES; attempt++)); do
        echo "Attempt $attempt/$MAX_RETRIES: $description"
        
        # Use timeout command to set timeout
        if timeout $timeout bash -c "$cmd"; then
            echo "$description completed successfully"
            return 0
        else
            local exit_code=$?
            echo "Attempt $attempt failed, exit code: $exit_code"
            
            if [ $attempt -lt $MAX_RETRIES ]; then
                echo "Waiting ${RETRY_DELAY} seconds before retry..."
                sleep $RETRY_DELAY
            else
                echo "$description failed, reached maximum retry attempts"
                return 1
            fi
        fi
    done
}

# Download function with retry
download_with_retry() {
    local url="$1"
    local output="$2"
    local description="$3"
    
    echo "Starting download: $description"
    
    for ((attempt=1; attempt<=MAX_RETRIES; attempt++)); do
        echo "Download attempt $attempt/$MAX_RETRIES: $description"
        
        if wget -O "$output" -c "$url"; then
            echo "$description downloaded successfully"
            return 0
        else
            echo "Download attempt $attempt failed"
            
            # Clean up potentially corrupted file
            if [ -f "$output" ]; then
                rm -f "$output"
            fi
            
            if [ $attempt -lt $MAX_RETRIES ]; then
                echo "Waiting ${RETRY_DELAY} seconds before retry download..."
                sleep $RETRY_DELAY
            else
                echo "$description download failed, reached maximum retry attempts"
                return 1
            fi
        fi
    done
}

# Function to import and tag Docker image with English messages
import_and_tag_image() {
    local tar_file="$1"
    local new_tag="$2"
    
    echo "========================================"
    echo "Docker Image Import and Tag Script"
    echo "========================================"
    
    # Check if tar file exists
    if [ ! -f "$tar_file" ]; then
        echo "Error: File $tar_file not found!"
        return 1
    fi
    
    # Import the Docker image
    echo "Step 1: Importing Docker image from $tar_file"
    local load_output=$(docker load -i "$tar_file" 2>&1)
    echo "Load command output: $load_output"
    
    # Extract image ID from output
    echo -e "\nStep 2: Extracting image ID from output"
    local image_id=""
    
    # 方法1: 从输出中提取镜像名，然后查询镜像ID
    if [[ $load_output =~ Loaded[[:space:]]image:[[:space:]]*([^[:space:]]+) ]]; then
        local image_name="${BASH_REMATCH[1]}"
        echo "Extracted image name from output: $image_name"
        
        # 等待一下确保镜像完全加载
        sleep 2
        
        # 使用镜像名查询镜像ID
        image_id=$(docker images -q "$image_name" 2>/dev/null)
        
        if [ -n "$image_id" ]; then
            echo "Found image ID via image name '$image_name': $image_id"
        fi
    fi
    
    # 方法2: 如果方法1失败，尝试从docker load输出中直接提取SHA256
    if [ -z "$image_id" ] && [[ $load_output =~ sha256:([a-f0-9]{64}) ]]; then
        local sha256_value="${BASH_REMATCH[1]}"
        image_id="${sha256_value:0:12}"
        echo "Found SHA256 pattern: sha256:${sha256_value}"
        echo "Using truncated SHA256 as image ID: $image_id"
    fi
    
    # 方法3: 如果前两种方法都失败，查找最近创建的镜像
    if [ -z "$image_id" ]; then
        echo "Warning: Could not extract image ID from output directly"
        echo "Attempting to find most recently loaded image..."
        
        # 查找所有镜像，按创建时间排序，取第一个
        image_id=$(docker images --quiet --filter "dangling=false" --format "{{.ID}}\t{{.CreatedAt}}" | sort -k2 -r | head -1 | cut -f1)
        
        if [ -n "$image_id" ]; then
            echo "Found most recent image: $image_id"
        else
            # 最后尝试：查找任何可用的镜像
            image_id=$(docker images --quiet | head -1)
            if [ -n "$image_id" ]; then
                echo "Found available image: $image_id"
            fi
        fi
    fi
    
    if [ -z "$image_id" ]; then
        echo "Error: Failed to determine image ID"
        return 1
    fi
    
    echo "Image ID determined: $image_id"
    
    # 验证镜像ID是否有效
    if ! docker inspect "$image_id" &>/dev/null; then
        echo "Error: Image ID '$image_id' is not valid"
        return 1
    fi
    
    # Tag the image
    echo -e "\nStep 3: Tagging image with new name: $new_tag"
    if docker tag "$image_id" "$new_tag"; then
        echo "SUCCESS: Image tagged successfully!"
        echo "New image reference: $new_tag"
        
        # Display result
        echo -e "\nStep 4: Verification"
        echo "Listing images with tag '$new_tag':"
        docker images | grep "$(echo "$new_tag" | cut -d: -f1)"
        
        return 0
    else
        echo "ERROR: Failed to tag image"
        return 1
    fi
}

# Export function to make it available in subshells
export -f import_and_tag_image

# Error handling function
handle_error() {
    local exit_code=$1
    local error_msg="$2"
    
    echo "ERROR: $error_msg (exit code: $exit_code)"
    echo "Installation process terminated due to error"
    exit $exit_code
}

# Set error handling
set -e
trap 'handle_error $? "Error occurred during script execution"' ERR

##########################################################################################
#  IrisNode installation script

echo "Starting IrisNode installation..."

# Temporarily disable IPv6
execute_with_retry "sysctl -w net.ipv6.conf.all.disable_ipv6=1 && sysctl -w net.ipv6.conf.default.disable_ipv6=1" "Disable IPv6" 30
sleep 2

# 禁用 UFW 防火墙
execute_with_retry "ufw disable" "Disable UFW firewall" 30

# 启用并启动 nftables
execute_with_retry "systemctl enable --now nftables" "Enable and start nftables" 30

# Switch to Aliyun source
echo "Shifting sources list to aliyun..."
# Backup source file
execute_with_retry "cp /etc/apt/sources.list /etc/apt/sources.list.backup" "Backup sources list" 30

# Output Aliyun source configuration to sources.list file
execute_with_retry 'tee /etc/apt/sources.list <<-EOF
deb https://mirrors.aliyun.com/ubuntu/ jammy main restricted universe multiverse
deb-src https://mirrors.aliyun.com/ubuntu/ jammy main restricted universe multiverse

deb https://mirrors.aliyun.com/ubuntu/ jammy-security main restricted universe multiverse
deb-src https://mirrors.aliyun.com/ubuntu/ jammy-security main restricted universe multiverse

deb https://mirrors.aliyun.com/ubuntu/ jammy-updates main restricted universe multiverse
deb-src https://mirrors.aliyun.com/ubuntu/ jammy-updates main restricted universe multiverse

# deb https://mirrors.aliyun.com/ubuntu/ jammy-proposed main restricted universe multiverse
# deb-src https://mirrors.aliyun.com/ubuntu/ jammy-proposed main restricted universe multiverse

deb https://mirrors.aliyun.com/ubuntu/ jammy-backports main restricted universe multiverse
deb-src https://mirrors.aliyun.com/ubuntu/ jammy-backports main restricted universe multiverse
EOF' "Configure Aliyun sources" 30

# Check br_netfilter
# Define module name to check
MODULE="br_netfilter"

# Check if module is included in /etc/modules
if ! grep -q "^$MODULE" /etc/modules; then
    echo "$MODULE not found in /etc/modules..."
    
    # Load module
    execute_with_retry "modprobe br_netfilter" "Load br_netfilter module" 30

    # Add br_netfilter to startup
    execute_with_retry 'echo "br_netfilter" >> /etc/modules' "Add br_netfilter to startup" 30
    echo "br_netfilter module added to /etc/modules"
else
    echo "$MODULE found in /etc/modules."
fi

# Check if Docker is installed via apt
if dpkg -l | grep -q docker-ce; then
    echo "Docker has been installed with apt."
else
    echo "Docker is not installed through apt, checking if there is a snap version of Docker..."

    # Check if snap version of Docker is installed
    if snap list | grep -q docker; then
        echo "Snap version of Docker has been installed and is uninstalling..."
        execute_with_retry "snap remove docker" "Uninstall snap Docker" 60
    fi 
    echo "Installing Docker through apt..."
    
    # Update apt package list
    execute_with_retry "apt update" "Update apt package list" 120
    
    # Install necessary packages
    execute_with_retry "apt install -y apt-transport-https ca-certificates curl wget software-properties-common openvswitch-switch unzip net-tools gnupg jq" "Install required dependency packages" 300

    #############################################################################################################################
    ## Online installation of docker-ce
    #############################################################################################################################
    echo "Installing Docker through apt..."

    execute_with_retry "install -m 0755 -d /etc/apt/keyrings" "Create apt keyrings directory" 30

    # Add Docker's GPG key
    execute_with_retry "curl -fsSL https://mirrors.aliyun.com/docker-ce/linux/ubuntu/gpg | gpg --batch --yes --dearmor -o /etc/apt/keyrings/docker.gpg" "Add Docker GPG key" 60
    execute_with_retry "chmod a+r /etc/apt/keyrings/docker.gpg" "Set Docker GPG key permissions" 30

    # Add Docker's official APT repository
    execute_with_retry 'echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://mirrors.aliyun.com/docker-ce/linux/ubuntu $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null' "Add Docker APT repository" 30

    # Update apt package list
    execute_with_retry "apt update" "Update Docker package list" 120
    
    # Install Docker
    execute_with_retry "apt install docker-ce docker-ce-cli containerd.io docker-buildx-plugin -y" "Install Docker components" 300

    echo "Docker has been installed with apt successfully."
fi

# Check if /etc/docker/ directory exists, create if not
if [ ! -d /etc/docker ]; then
    echo "/etc/docker directory does not exist, creating..."
    execute_with_retry "mkdir -p /etc/docker" "Create /etc/docker directory" 30
else
    echo "/etc/docker directory already exists."
fi

# Check if /etc/docker/daemon.json file exists
DAEMON_JSON_PATH="/etc/docker/daemon.json"
if [ -f "$DAEMON_JSON_PATH" ]; then
    echo "$DAEMON_JSON_PATH file already exists, backing up..."
    execute_with_retry "cp \"$DAEMON_JSON_PATH\" \"${DAEMON_JSON_PATH}.bak\"" "Backup Docker daemon configuration" 30
else
    echo "$DAEMON_JSON_PATH file does not exist, will write configuration directly..."
fi

# Write Docker image source configuration
echo "Writing Docker image source configuration..."

execute_with_retry 'tee /etc/docker/daemon.json <<-EOF
{
    "registry-mirrors": [
        "https://docker.hpcloud.cloud",
        "https://docker.m.daocloud.io",
        "https://docker.unsee.tech",
        "https://docker.1panel.live",
        "http://mirrors.ustc.edu.cn",
        "http://mirror.azure.cn",
        "https://dockerpull.org",
        "https://dockerhub.icu",
        "https://hub.rat.dev"
    ]
}
EOF' "Configure Docker image sources" 30

echo "Docker image source configuration has been written to $DAEMON_JSON_PATH."
echo "Restarting docker..."
execute_with_retry "systemctl daemon-reload && systemctl restart docker" "Restart Docker service" 60
echo "Docker restart complete."

# Check irisnode files
# Check if /etc/irisnode directory exists
IRISNODE_DIRECTORY="/opt/irisnode"
# Check if tar is installed
if ! command -v tar &> /dev/null; then
    echo "tar not installed, attempting to install..."
    execute_with_retry "apt update -qq && apt install -y tar" "Install tar tool" 120
    if [ $? -eq 0 ]; then
        echo "tar installed successfully."
    else
        echo "tar installation failed, please check system network or permissions." >&2
        exit 1
    fi
else
    echo "tar is already installed."
fi

# Check if /tmp/irisnode.tar.gz file exists
IRISNODE_TAR_GZ_FILE_PATH="/tmp/irisnode.tar.gz"
if [ -f "$IRISNODE_TAR_GZ_FILE_PATH" ]; then
    echo "$IRISNODE_TAR_GZ_FILE_PATH file already exists, removing and downloading newest version."
    rm "$IRISNODE_TAR_GZ_FILE_PATH" 
else
    echo "$IRISNODE_TAR_GZ_FILE_PATH file does not exist, will download..."
fi

cd /opt/

if [ -d "$IRISNODE_DIRECTORY" ]; then
    echo "Directory $IRISNODE_DIRECTORY already exists."
    # Delete irisnode folder
    execute_with_retry "rm -rf irisnode" "Remove old irisnode directory" 30
else
    echo "Directory $IRISNODE_DIRECTORY does not exist, downloading irisnode files..."
fi

download_with_retry "https://download.iris-cloud.cn/irisnode/$VERSION/irisnode.tar.gz" "/tmp/irisnode.tar.gz" "IrisNode installation package"
execute_with_retry "tar -xzvf /tmp/irisnode.tar.gz -C /opt/" "Extract IrisNode installation package" 60

# Check if jq tool is installed, exit if not
if ! command -v jq; then
    echo "jq tool is not installed. Please install it first."
    exit 1
fi &>/dev/null

# Try to read switch.json file content to variable, exit if failed
json_content=$(cat /opt/irisnode/etc/irisnode/switch/switch.json 2>/dev/null)
if [ -z "$json_content" ]; then
    echo "Failed to read the content of switch.json. Please check the file exists and has read permission."
    exit 1
fi

# Generate new random secret
new_secret=$(generate_random_secret)

# Use jq tool to directly replace the secret field in crypt field (ensure direct replacement of existing field)
new_json=$(echo "$json_content" | jq --arg new_secret "$new_secret" '.crypt.secret |= $new_secret')
if [ -z "$new_json" ]; then
    echo "Failed to update the 'crypt.secret' field using jq. Please check the JSON structure and the jq command."
    exit 1
fi

# Write modified content back to switch.json file, show error if failed
if ! echo "$new_json" > /opt/irisnode/etc/irisnode/switch/switch.json; then
    echo "Failed to write the modified content back to switch.json. Please check the file has write permission."
    exit 1
fi

echo "The 'crypt.secret' field in switch.json has been successfully updated."

cd $IRISNODE_DIRECTORY

################################################################################################################################
## Check if docker-compose is installed
################################################################################################################################
if command -v docker-compose > /dev/null 2>&1; then
    echo "docker-compose is already installed"
else
    echo "docker-compose is not installed, downloading and installing..."

    # Install docker-compose

    #####################################################################################################################################################
    ### Download offline version
    #####################################################################################################################################################
    download_with_retry "https://download.iris-cloud.cn/thirds/docker-compose/v2.34.0/docker-compose-linux-x86_64" "/usr/local/bin/docker-compose" "docker-compose tool"

    # Grant execution permissions
    execute_with_retry "chmod +x /usr/local/bin/docker-compose" "Set docker-compose execution permissions" 30

    # Verify installation
    if command -v docker-compose > /dev/null 2>&1; then
        echo "docker-compose installed successfully."
    else
        echo "docker-compose installation failed."
        exit 1
    fi
fi

# Start checking images
#declare -A IMAGES

# Check if image exists (name and version)
function check_image_exists() {
    local IMAGE_NAME=$1
    local IMAGE_VERSION=$2
    docker images --format "{{.Repository}}:{{.Tag}}" | grep -q "^${IMAGE_NAME}:${IMAGE_VERSION}$"
}

# Check if luscis/irisnode image exists
if check_image_exists "$IRISNODE_IMAGE" "$IRISNODE_TAG"; then
    echo "Image $IRISNODE_IMAGE:$IRISNODE_TAG already exists."
else
    # If luscis/irisnode image doesn't exist
    echo "Image luscis/irisnode:${IMAGES["luscis/irisnode"]} does not exist, downloading..."
    # Download offline image package
    download_with_retry "https://download.iris-cloud.cn/irisnode/$VERSION/images/image.x86_64.deb.tar.gz" "image.x86_64.deb.tar.gz" "IrisNode Docker image"
    
    # 使用修复后的导入函数
    execute_with_retry "import_and_tag_image \"image.x86_64.deb.tar.gz\" \"$IRISNODE_IMAGE:$IRISNODE_TAG\"" "Load IrisNode Docker image" 180
fi

# 注释掉所有 node-exporter 相关的代码，包括 prom-node-exporter-v1.6.1.tar 的下载
# if check_image_exists "$NODE_EXPORTER_IMAGE" "$NODE_EXPORTER_TAG"; then
#     echo "Image $NODE_EXPORTER_IMAGE:$NODE_EXPORTER_TAG already exists."
# else
#     # If prom/node_exporter image doesn't exist
#     echo "Image prom/node_exporter:${IMAGES["prom/node_exporter"]} does not exist, downloading..."
#     # Download offline image package
#     download_with_retry "https://download.iris-cloud.cn/irisnode/$VERSION/images/prom-node-exporter-v1.6.1.tar" "prom-node-exporter-v1.6.1.tar" "Node Exporter Docker image"
#     
#     # 使用修复后的导入函数
#     execute_with_retry "import_and_tag_image \"prom-node-exporter-v1.6.1.tar\" \"$NODE_EXPORTER_IMAGE:$NODE_EXPORTER_TAG\"" "Load Node Exporter Docker image" 180
# fi

# 注释掉所有 exporter 相关的 SSL 文件和配置下载
# # Download SSL files and configuration for exporter
# # Check if /root/exporter directory exists
# EXPORTER_DIRECTORY="/root/exporter"
# if [ -d "$EXPORTER_DIRECTORY" ]; then
#     echo "Directory $EXPORTER_DIRECTORY already exists."
#     # Delete /root/exporter folder
#     execute_with_retry "rm -rf /root/exporter" "Remove old exporter directory" 30
# else
#     echo "Directory $EXPORTER_DIRECTORY does not exist, downloading irisnode files..."
# fi
# 
# execute_with_retry "mkdir -p /root/exporter" "Create exporter directory" 30
# execute_with_retry "mkdir -p /root/exporter/ssl" "Create exporter ssl directory" 30
# download_with_retry "https://download.iris-cloud.cn/irisnode/$VERSION/config/config.yaml" "/root/exporter/config.yaml" "exporter configuration file"
# download_with_retry "https://download.iris-cloud.cn/irisnode/$VERSION/ssl/node_exporter.crt" "/root/exporter/ssl/node_exporter.crt" "exporter SSL certificate"
# download_with_retry "https://download.iris-cloud.cn/irisnode/$VERSION/ssl/node_exporter.key" "/root/exporter/ssl/node_exporter.key" "exporter SSL key"

# Check docker-compose configuration file
# Check if /home/node/docker-compose.yml file exists
DOCKER_COMPOSE_PATH="/opt/irisnode/docker-compose.yml"
if [ -f "$DOCKER_COMPOSE_PATH" ]; then
    echo "$DOCKER_COMPOSE_PATH file already exists, backing up..."
    execute_with_retry "cp \"$DOCKER_COMPOSE_PATH\" \"${DOCKER_COMPOSE_PATH}.bak\"" "Backup docker-compose configuration" 30
else
    echo "$DOCKER_COMPOSE_PATH file does not exist, will write configuration directly..."
fi

echo "Downloading irisnode docker-compose.yml..."
download_with_retry "https://download.iris-cloud.cn/irisnode/$VERSION/docker-compose.yml" "docker-compose.yml" "IrisNode docker-compose configuration"

echo "IrisNode reloading configuration file to start..."

# Check if container is running
container_name="irisnode-ipsec-1"

# Use docker ps to check if container is running
if docker ps --filter "name=$container_name" --filter "status=running" | grep "$container_name" > /dev/null; then
    echo "Stopping IrisNode Services..."
    # Stop all node-related containers
    execute_with_retry "docker-compose -f docker-compose.yml down" "Stop IrisNode services" 60
else
    echo "Starting IrisNode Services..."
fi

execute_with_retry "docker-compose -f docker-compose.yml up -d" "Start IrisNode services" 120
echo "IrisNode startup is complete."

# Final step: Replace certificate file in irisnode-switch-1 container and restart it
echo "Starting final configuration step: replacing certificate file in irisnode-switch-1 container..."

# Download the certificate file
download_with_retry "https://download.iris-cloud.cn/irisnode/$VERSION/crt.txt" "./crt.txt" "certificate file"

# Check if the container exists and is running
if docker ps --filter "name=irisnode-switch-1" --filter "status=running" | grep "irisnode-switch-1" > /dev/null; then
    echo "Container irisnode-switch-1 is running, proceeding with file replacement..."
    
    # Copy the certificate file to the container
    execute_with_retry "docker cp ./crt.txt irisnode-switch-1:/var/openlan/cert/crt" "Copy certificate file to container" 30
    
    # Restart the container
    execute_with_retry "docker restart irisnode-switch-1" "Restart irisnode-switch-1 container" 60
    
    # Clean up the downloaded certificate file
    rm -f ./crt.txt
    echo "Certificate file cleaned up"
    
    echo "Certificate replacement and container restart completed successfully"
else
    echo "ERROR: Container irisnode-switch-1 is not running. Cannot proceed with certificate replacement."
    exit 1
fi

# ===================== 新增清理逻辑 =====================
echo "Cleaning up prom/node-exporter:v1.6.1 image..."
# 停止/删除所有关联容器（即使没启动，防止残留）
docker stop $(docker ps -aq --filter "ancestor=prom/node-exporter:v1.6.1" 2>/dev/null) || true
docker rm $(docker ps -aq --filter "ancestor=prom/node-exporter:v1.6.1" 2>/dev/null) || true
# 删除镜像（包括虚悬镜像）
docker rmi prom/node-exporter:v1.6.1 2>/dev/null || true
docker rmi $(docker images --filter "reference=prom/node-exporter" -q) 2>/dev/null || true
echo "Cleanup completed."
# ===================== 新增结束 =====================

echo "IrisNode installation and configuration completed successfully!"