#!/bin/bash
SERVER_NAME=""
START_PORT=""
END_PORT=""

while [[ $# -gt 0 ]]; do
    case "$1" in
        --start-port)
            START_PORT="$2"
            shift 2
            ;;
        --end-port)
            END_PORT="$2"
            shift 2
            ;;
        --server)
            SERVER_NAME="$2"
            shift 2
            ;;
        -h|--help)
            echo "Usage: $0 --server hk/ca/or/deu --start-port xxxx --end-port xxxx"
            echo "Example:"
            echo "  $0 --server hk --start-port 40001 --end-port 40299"
            exit 0
            ;;
        *)
            echo "未知参数: $1"
            echo "使用 --help 查看帮助"
            exit 1
            ;;
    esac
done

if [ "$EUID" -ne 0 ]; then
    echo "请以 root 用户运行此脚本"
    exit 1
fi
case "$SERVER_NAME" in
  "hk")
    PORT_SERVER=16.162.219.193;;
  "alihk")
    PORT_SERVER=47.57.9.191;;
  "ca")
    PORT_SERVER=13.56.204.87;;
  "or")
    PORT_SERVER=52.89.207.198;;
  "deu")
    PORT_SERVER=2.56.97.114;;
  *)
    echo "未知的端口服务器"
    exit 1
    ;;
esac

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

PORT_FORWARD_SCRIPT="port_forward.sh"
SSH_KEY_FILE="remote.pem"
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
# 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
python3 /var/lib/vastai_kaalia/send_mach_info.py > /dev/null
echo "配置完成"