#!/bin/bash
# Info: Script to install Polkadot Zombienet
# Author: Provisioning System
# Release: 1.0
# Date: 2025-07-24

USAGE="install-polkadot-zombienet.sh"
[ "$1" == "-h" ] && echo "$USAGE" && exit 1

[ -r "env-polkadot-zombienet" ] && . ./env-polkadot-zombienet

ZOMBIENET_VERSION=${ZOMBIENET_VERSION:-1.3.133}
ZOMBIENET_PROVIDER=${ZOMBIENET_PROVIDER:-native}

# Determine architecture and OS
ARCH="$(uname -m)"
OS="$(uname -s | tr '[:upper:]' '[:lower:]')"

case "$ARCH" in
    x86_64) ZOMBIE_ARCH="x64" ;;
    aarch64|arm64) ZOMBIE_ARCH="arm64" ;;
    *) echo "Unsupported architecture: $ARCH" && exit 1 ;;
esac

case "$OS" in
    linux) ZOMBIE_OS="linux" ;;
    darwin) ZOMBIE_OS="macos" ;;
    *) echo "Unsupported OS: $OS" && exit 1 ;;
esac

ZOMBIENET_URL="https://github.com/paritytech/zombienet/releases/download"
ZOMBIENET_BINARY="zombienet-${ZOMBIE_OS}-${ZOMBIE_ARCH}"

ZOMBIENET_BIN_PATH=${ZOMBIENET_BIN_PATH:-/usr/local/bin}
ZOMBIENET_BINARY_NAME=${ZOMBIENET_BINARY:-zombienet}
ZOMBIENET_SYSTEMCTL_MODE=${ZOMBIENET_SYSTEMCTL_MODE:-enabled}

ZOMBIENET_CONFIG_PATH=${ZOMBIENET_CONFIG_PATH:-/etc/zombienet}
ZOMBIENET_WORK_PATH=${ZOMBIENET_WORK_PATH:-/var/lib/zombienet}
ZOMBIENET_NETWORKS_PATH=${ZOMBIENET_NETWORKS_PATH:-/var/lib/zombienet/networks}
ZOMBIENET_BINARIES_PATH=${ZOMBIENET_BINARIES_PATH:-/var/lib/zombienet/binaries}
ZOMBIENET_LOGS_PATH=${ZOMBIENET_LOGS_PATH:-/var/lib/zombienet/logs}

ZOMBIENET_RUN_USER=${ZOMBIENET_RUN_USER:-zombienet}
ZOMBIENET_RUN_GROUP=${ZOMBIENET_RUN_GROUP:-zombienet}
ZOMBIENET_RUN_USER_HOME=${ZOMBIENET_RUN_USER_HOME:-/home/zombienet}

echo "Installing Polkadot Zombienet ${ZOMBIENET_VERSION}..."

# Install dependencies based on provider
echo "Installing dependencies for provider: $ZOMBIENET_PROVIDER..."
if command -v apt-get >/dev/null 2>&1; then
    apt-get update
    DEPS="curl ca-certificates jq"
    
    case "$ZOMBIENET_PROVIDER" in
        "kubernetes")
            DEPS="$DEPS kubectl"
            ;;
        "podman")
            DEPS="$DEPS podman"
            ;;
        "native")
            # Native provider needs no additional system packages
            ;;
    esac
    
    apt-get install -y $DEPS
elif command -v yum >/dev/null 2>&1; then
    yum update -y
    DEPS="curl ca-certificates jq"
    
    case "$ZOMBIENET_PROVIDER" in
        "kubernetes")
            DEPS="$DEPS kubectl"
            ;;
        "podman")
            DEPS="$DEPS podman"
            ;;
    esac
    
    yum install -y $DEPS
elif command -v dnf >/dev/null 2>&1; then
    dnf update -y
    DEPS="curl ca-certificates jq"
    
    case "$ZOMBIENET_PROVIDER" in
        "kubernetes")
            DEPS="$DEPS kubectl"
            ;;
        "podman")
            DEPS="$DEPS podman"
            ;;
    esac
    
    dnf install -y $DEPS
else
    echo "Package manager not found. Please install dependencies manually."
    exit 1
fi

# Create user and group
if ! id "$ZOMBIENET_RUN_USER" &>/dev/null; then
    groupadd -r "$ZOMBIENET_RUN_GROUP"
    useradd -r -g "$ZOMBIENET_RUN_GROUP" -d "$ZOMBIENET_RUN_USER_HOME" -s /bin/bash -c "Zombienet service user" "$ZOMBIENET_RUN_USER"
fi

# Create directories
mkdir -p "$ZOMBIENET_CONFIG_PATH"
mkdir -p "$ZOMBIENET_WORK_PATH"
mkdir -p "$ZOMBIENET_NETWORKS_PATH"
mkdir -p "$ZOMBIENET_BINARIES_PATH"
mkdir -p "$ZOMBIENET_LOGS_PATH"
mkdir -p "$ZOMBIENET_RUN_USER_HOME"

# Download and install Zombienet binary
cd /tmp
echo "Downloading Zombienet from ${ZOMBIENET_URL}/v${ZOMBIENET_VERSION}/${ZOMBIENET_BINARY}..."
curl -L -o zombienet "${ZOMBIENET_URL}/v${ZOMBIENET_VERSION}/${ZOMBIENET_BINARY}"

if [ ! -f "zombienet" ]; then
    echo "Failed to download Zombienet binary"
    exit 1
fi

# Install binary
chmod +x zombienet
mv zombienet "$ZOMBIENET_BIN_PATH/$ZOMBIENET_BINARY_NAME"

# Download required binaries for native provider
if [ "$ZOMBIENET_PROVIDER" = "native" ]; then
    echo "Setting up binaries for native provider..."
    
    # Download Polkadot binary
    POLKADOT_VERSION="latest"
    POLKADOT_URL="https://github.com/paritytech/polkadot/releases/latest/download/polkadot"
    
    echo "Downloading Polkadot binary..."
    curl -L -o "$ZOMBIENET_BINARIES_PATH/polkadot" "$POLKADOT_URL"
    chmod +x "$ZOMBIENET_BINARIES_PATH/polkadot"
    
    # Download Polkadot-Parachain binary
    PARACHAIN_URL="https://github.com/paritytech/polkadot-sdk/releases/latest/download/polkadot-parachain"
    
    echo "Downloading Polkadot-Parachain binary..."
    curl -L -o "$ZOMBIENET_BINARIES_PATH/polkadot-parachain" "$PARACHAIN_URL" || {
        echo "Warning: Could not download polkadot-parachain binary"
        echo "You may need to build it manually or provide your own parachain binary"
    }
    
    if [ -f "$ZOMBIENET_BINARIES_PATH/polkadot-parachain" ]; then
        chmod +x "$ZOMBIENET_BINARIES_PATH/polkadot-parachain"
    fi
    
    # Create symbolic links in PATH
    ln -sf "$ZOMBIENET_BINARIES_PATH/polkadot" "$ZOMBIENET_BIN_PATH/polkadot"
    if [ -f "$ZOMBIENET_BINARIES_PATH/polkadot-parachain" ]; then
        ln -sf "$ZOMBIENET_BINARIES_PATH/polkadot-parachain" "$ZOMBIENET_BIN_PATH/polkadot-parachain"
    fi
fi

# Copy network templates
if [ -f "simple-network.toml" ]; then
    cp simple-network.toml "$ZOMBIENET_NETWORKS_PATH/"
fi

if [ -f "parachain-network.toml" ]; then
    cp parachain-network.toml "$ZOMBIENET_NETWORKS_PATH/"
fi

# Create default network configuration from template if available
if [ -f "network-config.toml.j2" ]; then
    cp network-config.toml.j2 "$ZOMBIENET_CONFIG_PATH/"
fi

# Set ownership
chown -R "$ZOMBIENET_RUN_USER:$ZOMBIENET_RUN_GROUP" "$ZOMBIENET_WORK_PATH"
chown -R "$ZOMBIENET_RUN_USER:$ZOMBIENET_RUN_GROUP" "$ZOMBIENET_NETWORKS_PATH"
chown -R "$ZOMBIENET_RUN_USER:$ZOMBIENET_RUN_GROUP" "$ZOMBIENET_BINARIES_PATH"
chown -R "$ZOMBIENET_RUN_USER:$ZOMBIENET_RUN_GROUP" "$ZOMBIENET_LOGS_PATH"
chown -R "$ZOMBIENET_RUN_USER:$ZOMBIENET_RUN_GROUP" "$ZOMBIENET_RUN_USER_HOME"
chown -R "$ZOMBIENET_RUN_USER:$ZOMBIENET_RUN_GROUP" "$ZOMBIENET_CONFIG_PATH"

# Create zombienet management script
cat > "$ZOMBIENET_BIN_PATH/zombienet-manager" << EOF
#!/bin/bash
# Zombienet Network Manager

ZOMBIENET_BIN="$ZOMBIENET_BIN_PATH/$ZOMBIENET_BINARY_NAME"
NETWORKS_PATH="$ZOMBIENET_NETWORKS_PATH"
LOGS_PATH="$ZOMBIENET_LOGS_PATH"
PROVIDER="$ZOMBIENET_PROVIDER"

case "\$1" in
    "spawn")
        NETWORK_FILE="\${2:-\$NETWORKS_PATH/simple-network.toml}"
        if [ ! -f "\$NETWORK_FILE" ]; then
            echo "Network file not found: \$NETWORK_FILE"
            exit 1
        fi
        echo "Spawning network with \$NETWORK_FILE using \$PROVIDER provider..."
        sudo -u $ZOMBIENET_RUN_USER "\$ZOMBIENET_BIN" spawn --provider "\$PROVIDER" "\$NETWORK_FILE"
        ;;
    "test")
        TEST_FILE="\${2}"
        if [ ! -f "\$TEST_FILE" ]; then
            echo "Test file not found: \$TEST_FILE"
            exit 1
        fi
        echo "Running test: \$TEST_FILE"
        sudo -u $ZOMBIENET_RUN_USER "\$ZOMBIENET_BIN" test --provider "\$PROVIDER" "\$TEST_FILE"
        ;;
    "setup")
        echo "Setting up Zombienet binaries..."
        sudo -u $ZOMBIENET_RUN_USER "\$ZOMBIENET_BIN" setup
        ;;
    "list")
        echo "Available network configurations:"
        ls -la "\$NETWORKS_PATH"/*.toml 2>/dev/null || echo "No network configurations found"
        ;;
    "logs")
        echo "Recent Zombienet logs:"
        find "\$LOGS_PATH" -name "*.log" -type f -exec tail -n 20 {} + 2>/dev/null || echo "No logs found"
        ;;
    "clean")
        echo "Cleaning up Zombienet processes and logs..."
        pkill -f zombienet || echo "No zombienet processes found"
        rm -rf "\$LOGS_PATH"/*
        echo "Cleanup completed"
        ;;
    "help"|*)
        echo "Zombienet Network Manager"
        echo "Usage: \$0 [command] [options]"
        echo ""
        echo "Commands:"
        echo "  spawn [network.toml]  Spawn a network (default: simple-network.toml)"
        echo "  test [test.zndsl]     Run a test file"
        echo "  setup                 Setup required binaries"
        echo "  list                  List available network configurations"
        echo "  logs                  Show recent logs"
        echo "  clean                 Clean up processes and logs"
        echo "  help                  Show this help message"
        echo ""
        echo "Provider: \$PROVIDER"
        echo "Networks path: \$NETWORKS_PATH"
        echo "Logs path: \$LOGS_PATH"
        ;;
esac
EOF

chmod +x "$ZOMBIENET_BIN_PATH/zombienet-manager"

# Create zombienet info file
cat > "$ZOMBIENET_CONFIG_PATH/zombienet-info.json" << EOF
{
  "version": "$ZOMBIENET_VERSION",
  "provider": "$ZOMBIENET_PROVIDER",
  "binary_path": "$ZOMBIENET_BIN_PATH/$ZOMBIENET_BINARY_NAME",
  "networks_path": "$ZOMBIENET_NETWORKS_PATH",
  "binaries_path": "$ZOMBIENET_BINARIES_PATH",
  "logs_path": "$ZOMBIENET_LOGS_PATH",
  "user": "$ZOMBIENET_RUN_USER",
  "manager_script": "$ZOMBIENET_BIN_PATH/zombienet-manager",
  "templates": [
    "simple-network.toml",
    "parachain-network.toml"
  ]
}
EOF

echo "=========================================="
echo "Polkadot Zombienet installation completed!"
echo "=========================================="
echo "Version: $ZOMBIENET_VERSION"
echo "Provider: $ZOMBIENET_PROVIDER"
echo "Binary: $ZOMBIENET_BIN_PATH/$ZOMBIENET_BINARY_NAME"
echo "Manager: $ZOMBIENET_BIN_PATH/zombienet-manager"
echo ""
echo "Paths:"
echo "Networks: $ZOMBIENET_NETWORKS_PATH"
echo "Binaries: $ZOMBIENET_BINARIES_PATH"
echo "Logs: $ZOMBIENET_LOGS_PATH"
echo "Config: $ZOMBIENET_CONFIG_PATH"
echo ""
echo "Quick start:"
echo "# List available networks"
echo "$ZOMBIENET_BIN_PATH/zombienet-manager list"
echo ""
echo "# Spawn simple network"
echo "$ZOMBIENET_BIN_PATH/zombienet-manager spawn"
echo ""
echo "# Spawn network with parachains"
echo "$ZOMBIENET_BIN_PATH/zombienet-manager spawn $ZOMBIENET_NETWORKS_PATH/parachain-network.toml"
echo ""
echo "# View logs"
echo "$ZOMBIENET_BIN_PATH/zombienet-manager logs"
echo ""
echo "# Clean up"
echo "$ZOMBIENET_BIN_PATH/zombienet-manager clean"

if [ "$ZOMBIENET_PROVIDER" = "native" ]; then
    echo ""
    echo "Native provider binaries installed:"
    ls -la "$ZOMBIENET_BINARIES_PATH"
fi

# Test Zombienet installation
echo ""
echo "Testing Zombienet installation..."
if "$ZOMBIENET_BIN_PATH/$ZOMBIENET_BINARY_NAME" --version >/dev/null 2>&1; then
    echo "✅ Zombienet binary is working"
else
    echo "⚠️  Zombienet binary test failed"
fi

# Cleanup
cd /
rm -rf /tmp/zombienet

echo ""
echo "Installation completed! Use 'zombienet-manager help' for usage instructions."