#!/bin/bash

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

SSH_FRPS=39.107.137.58

while [[ $# -gt 0 ]]; do
    case "$1" in
        --name)
            NAME="$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

# add param check
if [ -z "$NAME" ] || [ -z "$PORT" ]; then
    echo "Usage: $0 --name <name> --port <ssh_frp_port> [--port-server <public_ip>] [--vast-install-cmd <cmd>]"
    echo "Example:"
    echo "  $0 --name myserver --port 40001 --port-server 52.89.207.198 --vast-install-cmd 'wget https://console.vast.ai/install -O install; sudo python3 install 5ed661e763f06967f7f63ee839002e7c2a352cd6c88fa026cf4e65dfbb2d2c68 --interactive; history -d $((HISTCMD-1));'"
    exit 1
fi



# configure apt source
#cp /etc/apt/sources.list /etc/apt/sources.list.bak
#sed -i 's/archive.ubuntu.com/mirrors.aliyun.com/g' /etc/apt/sources.list
#sed -i 's/security.ubuntu.com/mirrors.aliyun.com/g' /etc/apt/sources.list
#apt update

# install base packages.
apt install -y git curl jq net-tools openssh-server
# configure ssh-server
echo "PermitRootLogin yes" >> /etc/ssh/sshd_config
mkdir -p "/root/.ssh"
cat ssh.pub > "/root/.ssh/authorized_keys"
systemctl restart ssh

# install ssh frpc
#tar zxf frp.tar.gz
#cd frp && bash install.sh "$NAME" "$SSH_FRPS" "$PORT"

# install driver
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
apt install -y build-essential software-properties-common dkms ubuntu-drivers-common
add-apt-repository ppa:graphics-drivers/ppa -y
apt update
apt install nvidia-driver-580-open -y


# install docker
curl -fsSL https://get.docker.com -o get-docker.sh && sudo sh get-docker.sh
apt install -y nvidia-docker2
nvidia-ctk runtime configure --runtime=docker

# configure docker daemon.json

cat <<EOF > /etc/docker/daemon.json
{
  "insecure-registries": [
    "hi0h0om60rpqq82fbb.xuanyuan.run",
    "hi0h0om60rpqq82fbb.xuanyuan.dev",
    "docker.xuanyuan.me"
  ],
  "registry-mirrors": [
    "https://hi0h0om60rpqq82fbb.xuanyuan.run",
    "https://hi0h0om60rpqq82fbb.xuanyuan.dev",
    "https://docker.xuanyuan.me"
  ],
  "runtimes": {
    "nvidia": {
      "args": [],
      "path": "/var/lib/vastai_kaalia/latest/kaalia_docker_shim"
    }
  }
}
EOF
systemctl restart docker

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

# 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" ubuntu@$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