#!/bin/bash

# Story Node Installation and Management Script

# Features:
# 1. Animated text display for a visually appealing start
# 2. Multiple installation and management options:
#    - Full node installation from scratch (including automatic snapshot download)
#    - Consensus layer snapshot download and application
#    - Execution layer (Geth) snapshot download and application
#    - Both consensus and execution layer snapshot application
#    - Consensus node upgrade
#    - Geth node upgrade
# 3. Flexible port configuration with default options
# 4. Improved error handling and directory management
# 5. Step-by-step guided process
# 6. Automatic dependency installation
# 7. Support for custom node moniker

# Function to display the animated text
display_text() {
    if ! command -v gem &> /dev/null; then
        echo "RubyGems (gem) is not installed. Installing Ruby..."
        sudo apt-get update
        sudo apt-get install -y ruby-full || { echo "Failed to install Ruby"; exit 1; }
    fi

    if ! command -v figlet &> /dev/null; then
        echo "Installing figlet..."
        sudo apt-get install -y figlet || { echo "Failed to install figlet"; exit 1; }
    fi

    if ! command -v lolcat &> /dev/null; then
        echo "Installing lolcat..."
        sudo gem install lolcat || { echo "Failed to install lolcat"; exit 1; }
    fi

    figlet "Story 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. Consensus Snapshot"
    echo "3. Geth Snapshot"
    echo "4. Both Node Snapshots"
    echo "5. Consensus Node Upgrade"
    echo "6. Geth Node Upgrade"
    read -p "Enter your choice (1-6): " INSTALL_TYPE

    case $INSTALL_TYPE in
        1|2|3|4|5|6) ;;
        *) echo "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
        export MONIKER

        echo "Please enter the port prefix for your node (press Enter for default ports):"
        read -p "Port prefix: " STORY_PORT
        if [ -z "$STORY_PORT" ]; then
            echo "Using default ports."
            unset STORY_PORT
        else
            echo "export STORY_PORT=$STORY_PORT" >> $HOME/.bash_profile
            source $HOME/.bash_profile
        fi

        echo "Setting up your node, $MONIKER, with ${STORY_PORT:+port prefix $STORY_PORT}${STORY_PORT:-default ports}..."
    fi
}

# Function to install dependencies
install_dependencies() {
    sudo apt update && sudo apt upgrade -y
    sudo apt install curl git wget htop tmux build-essential jq make lz4 gcc unzip aria2 -y
}

# Function to install Go
install_go() {
    if ! command -v go &> /dev/null; then
        cd $HOME
        curl https://dl.google.com/go/go1.23.1.linux-amd64.tar.gz | sudo tar -C/usr/local -zxvf -
        echo 'export GOROOT=/usr/local/go' >> $HOME/.bash_profile
        echo 'export GOPATH=$HOME/go' >> $HOME/.bash_profile
        echo 'export GO111MODULE=on' >> $HOME/.bash_profile
        echo 'export PATH=$PATH:/usr/local/go/bin:$HOME/go/bin' >> $HOME/.bash_profile
        source $HOME/.bash_profile
    fi
}

# Function to download and install binaries
install_binaries() {
    if [ ! -f $HOME/go/bin/geth ]; then
        cd $HOME
        wget -O geth https://github.com/piplabs/story-geth/releases/download/v0.10.0/geth-linux-amd64
        chmod +x $HOME/geth
        mv $HOME/geth ~/go/bin/
    fi

    if [ ! -f $HOME/go/bin/story ]; then
        cd $HOME
        rm -rf story
        git clone https://github.com/piplabs/story
        cd story
        git checkout v0.12.1
        go build -o story ./client
        sudo mv ~/story/story ~/go/bin/
        cd $HOME
    fi
}

# Function to initialize the node
initialize_node() {
    if [ ! -d $HOME/.story/story ]; then
        echo "Initializing Story node..."
        mkdir -p $HOME/.story/story
        mkdir -p $HOME/.story/geth
        story init --moniker $MONIKER --network odyssey
        if [ $? -ne 0 ]; then
            echo "Error: Failed to initialize Story node. Exiting."
            exit 1
        fi
    else
        echo "Story node directory already exists. Skipping initialization."
    fi

    # Download Genesis and Addrbook File
    if [ ! -f $HOME/.story/story/config/genesis.json ]; then
        echo "Downloading genesis file..."
        wget -O $HOME/.story/story/config/genesis.json https://files4.blacknodes.net/story-testnet/genesis.json --inet4-only
    else
        echo "Genesis file already exists. Skipping download."
    fi

    if [ ! -f $HOME/.story/story/config/addrbook.json ]; then
        echo "Downloading addrbook file..."
        wget -O $HOME/.story/story/config/addrbook.json https://files4.blacknodes.net/story-testnet/addrbook.json --inet4-only
    else
        echo "Addrbook file already exists. Skipping download."
    fi

    # Use Live Peers
    if [ -f $HOME/.story/story/config/config.toml ]; then
        echo "Updating peers..."
        peers=$(curl -s https://services.blacknodes.net/Story-Testnet/data/peers.txt)
        sed -i.bak -e "s/^persistent_peers *=.*/persistent_peers = \"$peers\"/" $HOME/.story/story/config/config.toml
    else
        echo "Error: config.toml not found. Make sure the node is initialized properly."
        exit 1
    fi

    # Set port configuration
    if [ ! -z "$STORY_PORT" ]; then
        echo "Setting custom ports..."
        if [ -f $HOME/.story/story/config/story.toml ]; then
            sed -i.bak -e "s%:1317%:${STORY_PORT}317%g;
            s%:8080%:${STORY_PORT}080%g;
            s%:9090%:${STORY_PORT}090%g;
            s%:9091%:${STORY_PORT}091%g;
            s%:8545%:${STORY_PORT}545%g;
            s%:8546%:${STORY_PORT}546%g;
            s%:6065%:${STORY_PORT}065%g" $HOME/.story/story/config/story.toml
        else
            echo "Warning: story.toml not found. Custom port configuration for Story may be incomplete."
        fi

        if [ -f $HOME/.story/story/config/config.toml ]; then
            sed -i.bak -e "s%:26658%:${STORY_PORT}658%g;
            s%:26657%:${STORY_PORT}657%g;
            s%:6060%:${STORY_PORT}060%g;
            s%:26656%:${STORY_PORT}656%g;
            s%^external_address = \"\"%external_address = \"$(wget -qO- eth0.me):${STORY_PORT}656\"%;
            s%:26660%:${STORY_PORT}660%g" $HOME/.story/story/config/config.toml
        else
            echo "Warning: config.toml not found. Custom port configuration may be incomplete."
        fi
    fi
}

# Function to create service files
create_service_files() {
    if [ ! -f /etc/systemd/system/story-geth.service ]; then
        sudo tee /etc/systemd/system/story-geth.service > /dev/null <<EOF
[Unit]
Description=Story Geth daemon
After=network-online.target

[Service]
User=$USER
ExecStart=$(which geth) --odyssey --syncmode full --http --http.api eth,net,web3,engine --http.vhosts '*' --http.addr 0.0.0.0 --http.port ${STORY_PORT:-8}545 --authrpc.port ${STORY_PORT:-8}551 --ws --ws.api eth,web3,net,txpool --ws.addr 0.0.0.0 --ws.port ${STORY_PORT:-8}546
Restart=on-failure
RestartSec=3
LimitNOFILE=65535

[Install]
WantedBy=multi-user.target
EOF
    fi

    if [ ! -f /etc/systemd/system/story.service ]; then
        sudo tee /etc/systemd/system/story.service > /dev/null <<EOF
[Unit]
Description=Story Service
After=network.target

[Service]
User=$USER
WorkingDirectory=$HOME/.story/story
ExecStart=$(which story) run
Restart=on-failure
RestartSec=5
LimitNOFILE=65535

[Install]
WantedBy=multi-user.target
EOF
    fi
}

# Function to download and apply consensus snapshot
apply_consensus_snapshot() {
    STORY_BLOCK_HEIGHT=$(curl -s https://files4.blacknodes.net/story-testnet/latest_block_height.txt)
    STORY_SNAPSHOT_FILE="odyssey_${STORY_BLOCK_HEIGHT}.tar.lz4"

    if [ -f "$STORY_SNAPSHOT_FILE" ]; then
        echo "Deleting existing Story snapshot file..."
        rm "$STORY_SNAPSHOT_FILE"
    fi

    echo "Downloading Story snapshot..."
    aria2c -x5 -s4 https://files4.blacknodes.net/story-testnet/${STORY_SNAPSHOT_FILE}

    echo "Story snapshot decompression started, please wait."
    if [ -f $HOME/.story/story/data/priv_validator_state.json ]; then
        cp $HOME/.story/story/data/priv_validator_state.json $HOME/.story/story/priv_validator_state.json.backup
    fi
    rm -rf $HOME/.story/story/data
    mkdir -p $HOME/.story/story/data
    lz4 -c -d ${STORY_SNAPSHOT_FILE} | tar -x -C $HOME/.story/story/data
    if [ -f $HOME/.story/story/priv_validator_state.json.backup ]; then
        mv $HOME/.story/story/priv_validator_state.json.backup $HOME/.story/story/data/priv_validator_state.json
    fi
}

# Function to download and apply Geth snapshot
apply_geth_snapshot() {
    GETH_BLOCK_HEIGHT=$(curl -s https://files4.blacknodes.net/story-testnet/el_geth_latest_block_height.txt)
    GETH_SNAPSHOT_FILE="story-geth_${GETH_BLOCK_HEIGHT}.tar.lz4"

    if [ -f "$GETH_SNAPSHOT_FILE" ]; then
        echo "Deleting existing Geth snapshot file..."
        rm "$GETH_SNAPSHOT_FILE"
    fi

    echo "Downloading Geth snapshot..."
    aria2c -x5 -s4 https://files4.blacknodes.net/story-testnet/${GETH_SNAPSHOT_FILE}

    echo "Geth snapshot decompression started, please wait."
    rm -rf $HOME/.story/geth/odyssey/geth/chaindata
    mkdir -p $HOME/.story/geth/odyssey/geth/chaindata
    lz4 -c -d ${GETH_SNAPSHOT_FILE} | tar -x -C $HOME/.story/geth/odyssey/geth/chaindata
}

# Function to upgrade consensus node
upgrade_consensus_node() {
    echo "Upgrading consensus node..."
    cd $HOME
    rm -rf story
    git clone https://github.com/piplabs/story
    cd story
    git checkout v0.12.1
    go build -o story ./client
    sudo mv ~/story/story ~/go/bin/
    sudo systemctl restart story
}

# Function to upgrade Geth node
upgrade_geth_node() {
    echo "Upgrading Geth node..."
    cd $HOME
    wget -O geth https://github.com/piplabs/story-geth/releases/download/v0.10.0/geth-linux-amd64
    chmod +x $HOME/geth
    mv $HOME/geth ~/go/bin/
    sudo systemctl restart story-geth
}

# Main script execution starts here
clear
display_text
sleep 3

choose_installation_type
setup_node_config

case $INSTALL_TYPE in
    1)
        echo "Starting full node installation..."
        install_dependencies
        install_go
        install_binaries
        initialize_node
        create_service_files
        apply_consensus_snapshot
        apply_geth_snapshot
        sudo systemctl daemon-reload
        sudo systemctl enable story story-geth
        sudo systemctl start story story-geth
        ;;
    2)
        initialize_node
        apply_consensus_snapshot
        sudo systemctl restart story
        ;;
    3)
        initialize_node
        apply_geth_snapshot
        sudo systemctl restart story-geth
        ;;
    4)
        initialize_node
        apply_consensus_snapshot
        apply_geth_snapshot
        sudo systemctl restart story story-geth
        ;;
    5)
        upgrade_consensus_node
        ;;
    6)
        upgrade_geth_node
        ;;
esac

echo "Story node operation completed successfully!"
echo "To view the Story logs, use: sudo journalctl -u story -f"
echo "To view the Geth logs, use: sudo journalctl -u story-geth -f"
echo "Press Ctrl+C to exit if you are stuck in the log view"