#!/bin/bash

# Set strict error handling
set -eo pipefail

# Color codes
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

# Functions
log() {
    echo -e "${GREEN}[$(date +%Y-%m-%d\ %H:%M:%S)]${NC} $1"
}

error() {
    echo -e "${RED}[$(date +%Y-%m-%d\ %H:%M:%S)] Error: $1${NC}"
}

warning() {
    echo -e "${YELLOW}[$(date +%Y-%m-%d\ %H:%M:%S)] Warning: $1${NC}"
}

# Function to display the animated text
display_text() {
    if ! command -v figlet &> /dev/null; then
        sudo apt-get update && sudo apt-get install -y figlet lolcat ruby > /dev/null 2>&1
    fi
    figlet "Avail Node Manager By BlackNodes" | lolcat
}

# Function to prompt for installation type
choose_installation_type() {
    echo "Please choose an installation or management option:"
    echo "1. Node Installation from Scratch"
    echo "2. Apply Snapshot"
    echo "3. Node Upgrade"
    read -p "Enter your choice (1-3): " INSTALL_TYPE

    case $INSTALL_TYPE in
        1|2|3) ;;
        *) error "Invalid option. Exiting."; exit 1 ;;
    esac
}

# Function to set up node configuration
setup_node_config() {
    if [ "$INSTALL_TYPE" == "1" ]; then
        echo "Please enter your node's MONIKER (name):"
        read MONIKER
        if [ -z "$MONIKER" ]; then
            MONIKER="avail-node-$(head /dev/urandom | tr -dc A-Za-z0-9 | head -c 6)"
            warning "No moniker provided. Using generated name: $MONIKER"
        fi
        export MONIKER
    fi
}

# Function to install dependencies
install_dependencies() {
    log "Installing dependencies..."
    sudo apt -qq update
    sudo apt -qq install -y curl git jq lz4 aria2 wget > /dev/null 2>&1
}

# Function to get snapshot info
get_snapshot_info() {
    local height_url="https://files4.blacknodes.net/avail/latest_block_height.txt"
    BLOCK_HEIGHT=$(curl -sfL "$height_url")
    if [ -z "$BLOCK_HEIGHT" ]; then
        error "Failed to fetch block height"
        exit 1
    fi
    SNAP_NAME="avail_da_mainnet_${BLOCK_HEIGHT}.tar.lz4"
    SNAPSHOT_URL="https://files4.blacknodes.net/avail/${SNAP_NAME}"
    
    # Verify snapshot exists
    if ! curl --output /dev/null --silent --head --fail "$SNAPSHOT_URL"; then
        error "Snapshot file not found: $SNAPSHOT_URL"
        exit 1
    fi
}

# Function to download and install binary
install_binary() {
    log "Installing Avail binary..."
    local os_version=$(lsb_release -sr | tr -d .)
    local binary_url="https://github.com/availproject/avail/releases/download/v2.2.5.0/x86_64-ubuntu-${os_version}-avail-node.tar.gz"
    
    if ! curl -L "$binary_url" | sudo tar -xz -C /usr/local/bin; then
        error "Failed to download or extract binary"
        exit 1
    fi
    
    # Verify installation
    if ! command -v avail-node &> /dev/null; then
        error "Binary installation failed"
        exit 1
    fi
}

# Function to create service file
create_service_file() {
    log "Creating systemd service..."
    sudo tee /etc/systemd/system/avail.service > /dev/null << EOF
[Unit]
Description=Avail Node
After=network-online.target

[Service]
User=$USER
ExecStart=$(which avail-node) \
    --base-path $HOME/.avail/data/ \
    --chain mainnet \
    --name "$MONIKER" \
    --rpc-external \
    --rpc-cors all \
    --rpc-methods safe
Restart=on-failure
RestartSec=10
LimitNOFILE=65535

[Install]
WantedBy=multi-user.target
EOF

    sudo systemctl daemon-reload
    sudo systemctl enable avail.service
}

# Function to download and apply snapshot
apply_snapshot() {
    log "Getting latest snapshot information..."
    get_snapshot_info
    
    log "Creating data directory..."
    mkdir -p $HOME/.avail/data/chains/avail_da_mainnet/paritydb

    log "Downloading snapshot from height $BLOCK_HEIGHT..."
    if ! aria2c -x5 -s5 --summary-interval=5 "$SNAPSHOT_URL"; then
        error "Failed to download snapshot"
        exit 1
    fi

    log "Applying snapshot..."
    if ! lz4 -c -d "${SNAP_NAME}" | tar -x -C $HOME/.avail/data/chains/avail_da_mainnet/paritydb; then
        error "Failed to extract snapshot"
        rm -f "${SNAP_NAME}"
        exit 1
    fi

    log "Cleaning up..."
    rm -f "${SNAP_NAME}"
}

# Function to upgrade node
upgrade_node() {
    log "Stopping Avail service..."
    sudo systemctl stop avail
    
    log "Installing new binary..."
    install_binary
    
    log "Starting Avail service..."
    sudo systemctl start avail
}

# Error handling
cleanup() {
    if [ $? -ne 0 ]; then
        error "Script failed! Check the error messages above."
        # Cleanup any partial installations or downloads
        rm -f $HOME/avail_da_mainnet_*.tar.lz4
    fi
}
trap cleanup EXIT

# Main script execution
clear
display_text
sleep 2

choose_installation_type
setup_node_config

case $INSTALL_TYPE in
    1)
        log "Starting full node installation..."
        install_dependencies
        install_binary
        create_service_file
        apply_snapshot
        sudo systemctl start avail
        ;;
    2)
        log "Applying latest snapshot..."
        sudo systemctl stop avail
        apply_snapshot
        sudo systemctl start avail
        ;;
    3)
        log "Upgrading node..."
        upgrade_node
        ;;
esac

log "Checking service status..."
sleep 5
if systemctl is-active --quiet avail; then
    log "Avail node is running successfully!"
    log "To view the node logs, use: sudo journalctl -u avail -f"
else
    error "Avail service is not running. Please check logs: sudo journalctl -u avail -f"
    exit 1
fi
