#!/bin/bash
set -e

# Subnet 118 - Wallet Creation and Staking Script
# This script sets up Python 3.12, creates a wallet, and stakes to Subnet 118
# Usage: curl -fsSL https://subnet-118-dashboard.vercel.app/scripts/new_wallet.sh | sh
# The script auto-detects your OS (Debian/Ubuntu via apt-get, Arch via pacman, macOS via uname)
# and interactively prompts for wallet name

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Function to print colored output
print_info() {
    echo -e "${BLUE}[INFO]${NC} $1" >&2
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1" >&2
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1" >&2
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1" >&2
}

# Check for required environment variables
print_info ""
print_info "========================================"
print_info "  Subnet 118 Wallet Setup"
print_info "========================================"
print_info ""

# Detect operating system automatically
print_info "Detecting operating system..."
if command -v apt-get &> /dev/null; then
    OS="debian_ubuntu"
    print_success "Detected: Debian/Ubuntu-based distribution (apt-get found)"
elif command -v pacman &> /dev/null; then
    OS="archlinux"
    print_success "Detected: Arch Linux (pacman found)"
elif [[ "$OSTYPE" == "darwin"* ]]; then
    OS="macos"
    print_success "Detected: macOS (Darwin)"
else
    print_error "Unsupported operating system"
    print_error "Could not detect apt-get (Debian/Ubuntu), pacman (Arch), or Darwin (macOS)"
    exit 1
fi

# Cache sudo credentials once for Linux systems (avoids multiple password prompts)
if [[ "$OS" != "macos" ]]; then
    print_info ""
    print_info "Administrator privileges required for installing Python..."
    sudo -v || {
        print_error "Failed to obtain sudo privileges"
        exit 1
    }
    # Keep sudo credentials fresh in the background
    (while true; do sudo -n true; sleep 60; kill -0 "$$" || exit; done 2>/dev/null &)
fi

# Prompt for wallet name interactively
print_info ""
print_info "========================================"
print_info "  Wallet Configuration"
print_info "========================================"
print_info ""

while true; do
    read -rp "Enter your wallet name (alphanumeric and underscores only): " WALLET_NAME

    if [[ -z "$WALLET_NAME" ]]; then
        print_error "Wallet name cannot be empty. Please try again."
        continue
    fi

    if [[ ! "$WALLET_NAME" =~ ^[a-zA-Z0-9_]+$ ]]; then
        print_error "Invalid wallet name: $WALLET_NAME"
        print_error "Wallet name must be alphanumeric (letters, numbers, underscores only)"
        print_info "Please try again."
        continue
    fi

    break
done

COLDKEY_NAME="${WALLET_NAME}_ck"
HOTKEY_NAME="${WALLET_NAME}_hk"

print_success "Wallet name set to: $WALLET_NAME"
print_info "Coldkey: $COLDKEY_NAME"
print_info "Hotkey: $HOTKEY_NAME"

# Install Python 3.12 based on OS
print_info ""
print_info "Installing Python 3.12..."

if [[ "$OS" == "macos" ]]; then
    if ! python3.12 --version &> /dev/null; then
        print_info "Installing Python 3.12 via Homebrew..."
        if ! command -v brew &> /dev/null; then
            print_warning "Homebrew not found. Installing Homebrew..."
            /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)" || {
                print_error "Failed to install Homebrew"
                exit 1
            }
            # Add Homebrew to PATH for Apple Silicon Macs
            if [[ -f "/opt/homebrew/bin/brew" ]]; then
                eval "$(/opt/homebrew/bin/brew shellenv)"
            fi
        fi
        brew install python@3.12 || {
            print_error "Failed to install Python 3.12"
            exit 1
        }
    else
        print_success "Python 3.12 is already installed"
    fi
    PYTHON_CMD="python3.12"

elif [[ "$OS" == "debian_ubuntu" ]]; then
    if ! python3.12 --version &> /dev/null; then
        # Clean up any previous Python build attempts
        sudo rm -rf /tmp/Python-3.* 2>/dev/null || true
        print_info "Installing Python 3.12..."

        # Update package list
        sudo apt update || {
            print_error "Failed to update package list"
            exit 1
        }

        # Try default repos first (fast path for newer systems)
        print_info "Attempting installation from default repositories..."
        if sudo apt install -y python3.12 python3.12-venv python3.12-dev 2>/dev/null; then
            print_success "Python 3.12 installed from default repos"
        else
            # Fall back to building from source (guaranteed to work)
            print_info "Python 3.12 not in default repos. Building from source..."
            print_warning "This will take 3-5 minutes but ensures compatibility."

            # Install build dependencies
            sudo apt install -y build-essential libssl-dev zlib1g-dev \
                libbz2-dev libreadline-dev libsqlite3-dev curl git \
                libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev \
                libffi-dev liblzma-dev wget ca-certificates || {
                print_error "Failed to install build dependencies"
                exit 1
            }

            # Download Python 3.12.8
            cd /tmp
            wget -q https://www.python.org/ftp/python/3.12.8/Python-3.12.8.tgz || {
                print_error "Failed to download Python 3.12"
                exit 1
            }

            print_info "Extracting Python source..."
            tar -xzf Python-3.12.8.tgz || {
                print_error "Failed to extract Python archive"
                exit 1
            }

            cd Python-3.12.8

            print_info "Configuring Python build..."
            ./configure --enable-optimizations --with-ensurepip=install --prefix=/usr/local 2>&1 | grep -E "(checking|creating)" || true

            print_info "Building Python (this takes 3-5 minutes)..."
            make -j$(nproc) > /dev/null 2>&1 || {
                print_error "Failed to build Python"
                exit 1
            }

            sudo make altinstall > /dev/null 2>&1 || {
                print_error "Failed to install Python"
                exit 1
            }

            # Verify installation
            if ! /usr/local/bin/python3.12 --version &> /dev/null; then
                print_error "Python 3.12 installation verification failed"
                exit 1
            fi

            # Create symlink if needed
            if ! command -v python3.12 &> /dev/null; then
                sudo ln -sf /usr/local/bin/python3.12 /usr/bin/python3.12
            fi

            # Clean up
            cd /tmp
            rm -rf Python-3.12.8 Python-3.12.8.tgz

            print_success "Python 3.12 built and installed successfully"
        fi

        # Ensure pip is available for Python 3.12
        if ! python3.12 -m pip --version &> /dev/null; then
            sudo apt install -y python3-pip 2>/dev/null || true
        fi
    else
        print_success "Python 3.12 is already installed"
    fi
    PYTHON_CMD="python3.12"

elif [[ "$OS" == "archlinux" ]]; then
    if ! python3.12 --version &> /dev/null; then
        print_info "Installing Python 3.12 from official repos..."

        # Python is in official Arch repos, no AUR needed
        sudo pacman -Sy --noconfirm python python-pip || {
            print_error "Failed to install Python"
            exit 1
        }

        # Check if we got Python 3.12+
        PYTHON_VERSION=$(python --version 2>&1 | awk '{print $2}' | cut -d. -f1,2)

        if [[ "$PYTHON_VERSION" == "3.12" ]] || [[ "$PYTHON_VERSION" == "3.13" ]] || [[ "$PYTHON_VERSION" > "3.12" ]]; then
            print_success "Python $PYTHON_VERSION installed"
            # Create python3.12 symlink if needed
            if ! command -v python3.12 &> /dev/null; then
                sudo ln -sf /usr/bin/python /usr/bin/python3.12
            fi
        else
            print_error "Arch should have Python 3.12+. Found version: $PYTHON_VERSION"
            print_error "Please update your system: sudo pacman -Syu"
            exit 1
        fi
    else
        print_success "Python 3.12 is already installed"
    fi
    PYTHON_CMD="python3.12"
fi

# Verify Python installation
if ! $PYTHON_CMD --version &> /dev/null; then
    print_error "Python 3.12 installation verification failed"
    exit 1
fi

print_success "Python 3.12 ready: $($PYTHON_CMD --version)"

# Ensure pip is installed (but don't use ensurepip on Arch - it causes errors)
print_info "Ensuring pip is installed..."
if ! $PYTHON_CMD -m pip --version &> /dev/null; then
    if [[ "$OS" != "archlinux" ]]; then
        # Only try ensurepip on non-Arch systems
        print_info "Installing pip..."
        $PYTHON_CMD -m ensurepip --upgrade || {
            print_error "Failed to install pip"
            exit 1
        }
        print_success "pip installed"
    else
        # On Arch, pip should already be installed via python-pip package
        print_warning "pip not found but should be installed via python-pip package"
    fi
else
    print_success "pip is already installed"
fi

# Create working directory
WORK_DIR="$HOME/118_miner"
print_info ""
print_info "Creating working directory: $WORK_DIR"
mkdir -p "$WORK_DIR"
cd "$WORK_DIR"

# Create virtual environment
print_info "Creating Python virtual environment..."
$PYTHON_CMD -m venv .venv
print_success "Virtual environment created"

# Install bittensor packages
print_info ""
print_info "Installing bittensor and bittensor-cli..."
.venv/bin/pip install bittensor bittensor-cli
print_success "Packages installed"

# Create wallet
print_info ""
print_info "========================================"
print_info "  Creating Wallet"
print_info "========================================"
print_warning "IMPORTANT: You will be shown mnemonic seeds."
print_warning "WRITE THEM DOWN AND STORE THEM SECURELY!"
print_warning ""
print_warning "Press Enter to continue..."
read -r

print_info "Creating wallet..."
.venv/bin/btcli wallet create --wallet.name "$COLDKEY_NAME" --wallet.hotkey "$HOTKEY_NAME"

print_info ""
print_success "Wallet created successfully!"
print_info ""
print_warning "!!! SAVE YOUR COLDKEY AND HOTKEY MNEMONIC SEEDS !!!"
print_info ""

# Get coldkey address (parsing the actual btcli output format)
COLDKEY_OUTPUT=$(.venv/bin/btcli wallet list --wallet.name "$COLDKEY_NAME" 2>/dev/null)
# Method 1: Look for ss58_address field with regex
COLDKEY_ADDRESS=$(echo "$COLDKEY_OUTPUT" | grep -oP '(?<=ss58_address:\s)5[A-Za-z0-9]{47}' | head -1)

# Method 2: If that fails, try looking for any 5-prefixed base58 address (48 chars total)
if [[ -z "$COLDKEY_ADDRESS" ]]; then
    COLDKEY_ADDRESS=$(echo "$COLDKEY_OUTPUT" | grep -oE '5[A-Za-z0-9]{47}' | head -1)
fi

# Method 3: Try JSON output format
if [[ -z "$COLDKEY_ADDRESS" ]]; then
    COLDKEY_JSON=$(.venv/bin/btcli wallet list --wallet.name "$COLDKEY_NAME" --json 2>/dev/null || echo "")
    if [[ -n "$COLDKEY_JSON" ]]; then
        COLDKEY_ADDRESS=$(echo "$COLDKEY_JSON" | .venv/bin/python -c "
import sys, json
try:
    data = json.load(sys.stdin)
    # Try different possible JSON structures
    if isinstance(data, dict):
        # Look for coldkey or ss58_address fields
        for key in ['coldkey', 'ss58_address', 'address']:
            if key in data:
                print(data[key])
                break
        # Search nested structures
        for value in data.values():
            if isinstance(value, dict) and 'ss58_address' in value:
                print(value['ss58_address'])
                break
except:
    pass
" 2>/dev/null)
    fi
fi

# Method 4: Try inspect command as fallback
if [[ -z "$COLDKEY_ADDRESS" ]]; then
    INSPECT_OUTPUT=$(.venv/bin/btcli wallet inspect --wallet.name "$COLDKEY_NAME" 2>/dev/null || echo "")
    if [[ -n "$INSPECT_OUTPUT" ]]; then
        COLDKEY_ADDRESS=$(echo "$INSPECT_OUTPUT" | grep -oE '5[A-Za-z0-9]{47}' | head -1)
    fi
fi

if [[ -n "$COLDKEY_ADDRESS" ]]; then
    print_info ""
    echo "Your coldkey address: $COLDKEY_ADDRESS"
    print_info ""
else
    print_warning "Could not auto-detect coldkey address."
    print_info "Run this command to view your coldkey address:"
    print_info "  cd $WORK_DIR && .venv/bin/btcli wallet list --wallet.name $COLDKEY_NAME"
fi

print_info ""
print_info "========================================"
print_info "  Index Selection"
print_info "========================================"
print_info ""
print_info "Available indexes:"
print_info "  0 = TSBCSI"
print_info "  1 = Top 10"
print_info "  2 = Full Stack"
print_info "  3 = Fintech"
print_info "  4 = Bittensor Universe"
print_info ""
print_info "Learn more at: https://app.trustedstake.ai/strat"
print_info ""

# Prompt for index choice
while true; do
    read -rp "Which index do you wish to delegate to (0-4)? " INDEX_CHOICE

    if [[ "$INDEX_CHOICE" =~ ^[0-4]$ ]]; then
        break
    else
        print_error "Invalid choice. Please enter a number between 0 and 4."
    fi
done

# Set required balance based on index
if [[ "$INDEX_CHOICE" == "4" ]]; then
    REQUIRED_BALANCE=20.0
    print_info ""
    print_warning "Bittensor Universe requires 20 TAO minimum"
else
    REQUIRED_BALANCE=2.0
    print_info ""
    print_info "This index requires 2 TAO minimum"
fi

print_info ""
print_info "========================================"
print_info "  Balance Check"
print_info "========================================"
print_info ""
print_info "Please send $REQUIRED_BALANCE TAO (or more) to your coldkey address: $COLDKEY_ADDRESS"
print_info "The script will check for funds every 30 seconds."
print_info ""
print_info "Press Enter when you've sent the funds..."
read -r

print_info "Checking wallet balance..."
print_info ""

# Poll for balance
while true; do
    BALANCE_OUTPUT=$(.venv/bin/btcli w balance --wallet.name "$COLDKEY_NAME" --network finney --json-out 2>/dev/null || echo "")

    if [[ -n "$BALANCE_OUTPUT" ]]; then
        # Extract free balance from JSON output (use venv python, not system python3)
        FREE_BALANCE=$(echo "$BALANCE_OUTPUT" | .venv/bin/python -c "
import sys, json
try:
    data = json.load(sys.stdin)
    balances = data.get('balances', {})
    for wallet_name, wallet_data in balances.items():
        print(wallet_data.get('free', 0))
        break
except:
    print(0)
" 2>/dev/null || echo "0")

        # Check if we got a valid number
        if [[ -n "$FREE_BALANCE" ]] && [[ "$FREE_BALANCE" =~ ^[0-9]+\.?[0-9]*$ ]]; then
            # Use awk for float comparison (more portable than bc)
            if awk "BEGIN {exit !($FREE_BALANCE >= $REQUIRED_BALANCE)}"; then
                print_success "Balance sufficient: $FREE_BALANCE TAO"
                break
            else
                print_warning "Waiting for sufficient funds... (Current: $FREE_BALANCE TAO, Required: $REQUIRED_BALANCE TAO)"
            fi
        else
            print_warning "Waiting for funds... (Required: $REQUIRED_BALANCE TAO)"
        fi
    else
        print_warning "Unable to check balance. Retrying..."
    fi

    sleep 30
done

print_info ""
print_success "Balance check passed!"

# Registration and delegation
print_info ""
print_info "========================================"
print_info "  Registration & Delegation"
print_info "========================================"
print_info ""
print_info "Total one-time fees = Neuron registration cost (+transaction fees) + staking proxy transaction fees"
print_info "Locked funds: ~0.1 TAO locked as a proxy reserve (refunded on stake withdrawal)"
print_info ""

# Register neuron
print_info "Registering as a neuron on Subnet 118..."
.venv/bin/btcli subnets register --netuid 118 --wallet.name "$COLDKEY_NAME" --wallet.hotkey "$HOTKEY_NAME" --network finney

print_success "Registration complete!"
print_info ""

# Download and run miner_cli.py for delegation (use curl or wget depending on availability)
print_info "Downloading miner delegation script..."
mkdir -p "ETF/core"
if command -v curl &> /dev/null; then
    curl -fsSL -o ETF/core/constants.py https://raw.githubusercontent.com/mobiusfund/etf/refs/heads/main/ETF/core/constants.py
    curl -fsSL -o miner_cli.py https://raw.githubusercontent.com/mobiusfund/etf/refs/heads/main/contrib/miner_cli/miner_cli.py
elif command -v wget &> /dev/null; then
    wget -q -O ETF/core/constants.py https://raw.githubusercontent.com/mobiusfund/etf/refs/heads/main/ETF/core/constants.py
    wget -q -O miner_cli.py https://raw.githubusercontent.com/mobiusfund/etf/refs/heads/main/contrib/miner_cli/miner_cli.py
else
    print_error "Neither curl nor wget found. Please install one of them."
    exit 1
fi

print_info "Delegating to index $INDEX_CHOICE..."
.venv/bin/python miner_cli.py delegate \
    --index "$INDEX_CHOICE" \
    --wallet.name "$COLDKEY_NAME" \
    --wallet.hotkey "$HOTKEY_NAME" \
    --network finney \
    --yes

print_info ""
print_success "========================================"
print_success "  Setup Complete!"
print_success "========================================"
print_success ""
print_success "Your wallet is now set up and delegating to index $INDEX_CHOICE"
print_success ""
print_info "Working directory: $WORK_DIR"
print_info "To check your wallet: btcli wallet list --wallet.name $COLDKEY_NAME"
print_info "To view your miner on the dashboard: https://subnet-118-dashboard.vercel.app"
print_info ""
print_warning "Remember: Keep your mnemonic seeds safe!"