#!/bin/bash

NAME=""
PORT=""                 # ssh frp port
PORT_SERVER="52.89.207.198"
VAST_INSTALL_KEY=""
START_PORT=""
END_PORT=""
USER="ubuntu"

SSH_FRPS=39.107.137.58

while [[ $# -gt 0 ]]; do
    case "$1" in
        --name)
            NAME="$2"
            shift 2
            ;;
        --user)
            USER="$2"
            shift 2
            ;;
        --port)
            PORT="$2"
            shift 2
            ;;
        --start-port)
            START_PORT="$2"
            shift 2
            ;;
        --end-port)
            END_PORT="$2"
            shift 2
            ;;
        --public-ip)
            PORT_SERVER="$2"
            shift 2
            ;;
        --vast-install-key)
            VAST_INSTALL_KEY="$2"
            shift 2
            ;;
        -h|--help)
            echo "Usage: $0 --name <name> --port <ssh_frp_port> [--public-ip <public_ip>] [--vast-install-key <key>]"
            echo "Example:"
            echo "  $0 --name myserver --port 40001 --public-ip 52.89.207.198 --vast-install-key '5ed661e763f06967f7f63ee839002e7c2a352cd6c88fa026cf4e65dfbb2d2c68'"
            exit 0
            ;;
        *)
            echo "未知参数: $1"
            echo "使用 --help 查看帮助"
            exit 1
            ;;
    esac
done

if [ "$EUID" -ne 0 ]; then
    echo "请以 root 用户运行此脚本"
    exit 1
fi


# install base packages.
apt install -y git curl jq net-tools openssh-server


# install tailscale
curl -fsSL https://tailscale.com/install.sh | sh
sudo tailscale up --authkey tskey-auth-k8BXZ58bvH11CNTRL-RXkFujgcp2ihihZEJmVZ2iagaLg22aJg

# get tailscale ip with ip addr
tailscale_ip=$(ip addr | grep tailscale0 | grep inet | awk '{print $2}' | cut -d/ -f1)

# add speedtest
docker pull tscel/speedtest:800_1600
docker tag tscel/speedtest:800_1600 vastai/test:speedtest

# execute port_forward.sh on remote server.
PORT_FORWARD_SCRIPT="port_forward.sh"
SSH_KEY_FILE="remote.pem"

# ssh port is 30004, so start from 40001
#start_port=$(( (40 + $PORT-30004) * 1000 + 1 ))
#end_port=$(( (40 + $PORT-30004) * 1000 + 299 ))
PORT_RANGE="${START_PORT}-${END_PORT}"

if [ ! -f "$SSH_KEY_FILE" ]; then
  echo "错误: 未找到 SSH 密钥文件 $SSH_KEY_FILE，无法通过 remote.pem 连接远端服务器" >&2
  exit 1
fi

if [ -f "$PORT_FORWARD_SCRIPT" ]; then
  ssh -i "$SSH_KEY_FILE" -o "StrictHostKeyChecking=no" $USER@$PORT_SERVER \
    "sudo bash -s -- 'add' '$tailscale_ip' '$PORT_RANGE'" < "$PORT_FORWARD_SCRIPT"
else
  echo "警告: 未找到 $PORT_FORWARD_SCRIPT，无法在$PORT_SERVER 远端执行 port_forward.sh" >&2
fi

# then execute vast install cmd
if [ -n "$VAST_INSTALL_KEY" ];then
    python3 ./install $VAST_INSTALL_KEY --no-docker --interactive; history -d $((HISTCMD-1));
fi

# set ip and port.
bash -c "echo -n \"$PORT_SERVER\" > /var/lib/vastai_kaalia/host_ipaddr"
echo -n "$PORT_RANGE" > /var/lib/vastai_kaalia/host_port_range
systemctl restart vastai

# disable apt auto update.
systemctl disable --now unattended-upgrades