#!/bin/bash

RED='\033[0;31m'
GREEN='\033[0;32m'
CYAN='\033[0;36m'
YELLOW='\033[0;33m'
NC='\033[0m'

DEFAULT_INSTALL_DIR="/opt/realm"
INSTALL_DIR=""
REALM_BIN=""
REALM_CONF=""
REALM_LOG=""
REALM_SERVICE="/etc/systemd/system/realm.service"
SCRIPT_RAW_URL="https://git.suhang.me/suhang/scripts/raw/branch/release/realm/realm.sh"
GITHUB_API="https://api.github.com/repos/zhboner/realm/releases/latest"
DEFAULT_PROXY="ghproxy.com"
GITHUB_PROXY=""

function set_paths() {
    INSTALL_DIR="$1"
    REALM_BIN="${INSTALL_DIR}/realm"
    REALM_CONF="${INSTALL_DIR}/config.toml"
    REALM_LOG="${INSTALL_DIR}/realm.log"
}

function detect_install_dir() {
    local dir=""
    if [[ -f "$REALM_SERVICE" ]]; then
        dir=$(grep -oE 'ExecStart=[^ ]+' "$REALM_SERVICE" | head -n1 | sed -E 's|ExecStart=(.+)/realm$|\1|')
    fi
    [[ -z "$dir" ]] && dir="$DEFAULT_INSTALL_DIR"
    set_paths "$dir"
    load_proxy
}

function gh_url() {
    local url="$1"
    if [[ -n "$GITHUB_PROXY" ]]; then
        echo "https://${GITHUB_PROXY}/${url}"
    else
        echo "$url"
    fi
}

function load_proxy() {
    if [[ -n "$INSTALL_DIR" && -f "${INSTALL_DIR}/.proxy" ]]; then
        GITHUB_PROXY=$(head -n1 "${INSTALL_DIR}/.proxy" 2>/dev/null | tr -d '[:space:]')
    fi
}

function prompt_proxy() {
    local p
    read -p "请输入 GitHub 代理域名 [默认 ${DEFAULT_PROXY}，输入 none 不使用代理]: " p
    p="${p:-$DEFAULT_PROXY}"
    if [[ "$p" == "none" || "$p" == "NONE" ]]; then
        GITHUB_PROXY=""
        echo "GitHub proxy: disabled"
    else
        p="${p#http://}"
        p="${p#https://}"
        p="${p%/}"
        GITHUB_PROXY="$p"
        echo "GitHub proxy: ${GITHUB_PROXY}"
    fi
    mkdir -p "$INSTALL_DIR"
    printf '%s\n' "$GITHUB_PROXY" > "${INSTALL_DIR}/.proxy"
}

function prompt_install_dir() {
    local dir
    read -p "请输入安装目录 [默认 ${DEFAULT_INSTALL_DIR}]: " dir
    dir="${dir:-$DEFAULT_INSTALL_DIR}"
    if [[ "$dir" != /* ]]; then
        echo -e "${RED}必须是绝对路径。${NC}"
        return 1
    fi
    set_paths "$dir"
    mkdir -p "$INSTALL_DIR" || return 1
    echo "Install directory: ${INSTALL_DIR}"
}

function check_root() {
    if [[ $EUID -ne 0 ]]; then
        echo -e "${RED}This script must be run as root.${NC}"
        exit 1
    fi
}

function check_dependencies() {
    local deps=(curl tar systemctl)
    local missing=()
    for d in "${deps[@]}"; do
        if ! command -v "$d" >/dev/null 2>&1; then
            missing+=("$d")
        fi
    done
    if [[ ${#missing[@]} -gt 0 ]]; then
        echo -e "${RED}Missing dependencies: ${missing[*]}${NC}"
        echo -e "${YELLOW}Please install them and re-run the script.${NC}"
        exit 1
    fi
}

function detect_arch() {
    local m
    m=$(uname -m)
    case "$m" in
        x86_64|amd64)   echo "x86_64-unknown-linux-gnu" ;;
        aarch64|arm64)  echo "aarch64-unknown-linux-gnu" ;;
        armv7l|armv7)   echo "armv7-unknown-linux-gnueabihf" ;;
        *)
            echo -e "${RED}Unsupported architecture: $m${NC}" >&2
            return 1
            ;;
    esac
}

function get_latest_version() {
    local api_url v
    api_url=$(gh_url "$GITHUB_API")
    v=$(curl -fsSL "$api_url" 2>/dev/null | grep -oE '"tag_name":\s*"[^"]+"' | head -n1 | sed -E 's/.*"([^"]+)"$/\1/')
    if [[ -z "$v" ]]; then
        echo -e "${RED}Failed to fetch latest realm version from GitHub (proxy: ${GITHUB_PROXY:-none}).${NC}" >&2
        return 1
    fi
    echo "$v"
}

function download_realm() {
    local version="$1"
    local arch
    arch=$(detect_arch) || return 1

    local url
    url=$(gh_url "https://github.com/zhboner/realm/releases/download/${version}/realm-${arch}.tar.gz")
    local tmp
    tmp=$(mktemp -d)

    echo "Downloading realm ${version} (${arch}) via ${GITHUB_PROXY:-direct}..."
    if ! curl -fsSL "$url" -o "${tmp}/realm.tar.gz"; then
        echo -e "${RED}Download failed: $url${NC}"
        rm -rf "$tmp"
        return 1
    fi

    if ! tar -xzf "${tmp}/realm.tar.gz" -C "$tmp"; then
        echo -e "${RED}Failed to extract realm archive.${NC}"
        rm -rf "$tmp"
        return 1
    fi

    if [[ ! -f "${tmp}/realm" ]]; then
        echo -e "${RED}realm binary not found in archive.${NC}"
        rm -rf "$tmp"
        return 1
    fi

    install -m 0755 "${tmp}/realm" "$REALM_BIN"
    rm -rf "$tmp"
    echo -e "${GREEN}realm installed to ${REALM_BIN}${NC}"
}

function write_default_config() {
    mkdir -p "$INSTALL_DIR"
    if [[ -f "$REALM_CONF" ]]; then
        return 0
    fi
    cat > "$REALM_CONF" <<EOF
[log]
level = "warn"
output = "${REALM_LOG}"

[network]
no_tcp = false
use_udp = true

EOF
    echo "Default config written to ${REALM_CONF}"
}

function write_systemd_unit() {
    cat > "$REALM_SERVICE" <<EOF
[Unit]
Description=Realm TCP/UDP relay
Documentation=https://github.com/zhboner/realm
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
User=root
Restart=on-failure
RestartSec=5s
LimitNOFILE=1048576
ExecStart=${REALM_BIN} -c ${REALM_CONF}

[Install]
WantedBy=multi-user.target
EOF
    systemctl daemon-reload
    systemctl enable realm >/dev/null 2>&1
    echo "Systemd unit written to ${REALM_SERVICE}"
}

function open_firewall_port() {
    local port="$1"
    [[ -z "$port" ]] && return 0

    if command -v ufw >/dev/null 2>&1 && ufw status 2>/dev/null | grep -q "Status: active"; then
        ufw allow "${port}/tcp" >/dev/null 2>&1
        ufw allow "${port}/udp" >/dev/null 2>&1
        echo "ufw: opened ${port}/tcp,udp"
        return 0
    fi

    if command -v firewall-cmd >/dev/null 2>&1 && firewall-cmd --state 2>/dev/null | grep -q running; then
        firewall-cmd --zone=public --add-port="${port}/tcp" --permanent >/dev/null 2>&1
        firewall-cmd --zone=public --add-port="${port}/udp" --permanent >/dev/null 2>&1
        firewall-cmd --reload >/dev/null 2>&1
        echo "firewalld: opened ${port}/tcp,udp"
        return 0
    fi

    if command -v iptables >/dev/null 2>&1; then
        iptables -C INPUT -p tcp --dport "$port" -j ACCEPT 2>/dev/null || iptables -A INPUT -p tcp --dport "$port" -j ACCEPT
        iptables -C INPUT -p udp --dport "$port" -j ACCEPT 2>/dev/null || iptables -A INPUT -p udp --dport "$port" -j ACCEPT
        if command -v ip6tables >/dev/null 2>&1; then
            ip6tables -C INPUT -p tcp --dport "$port" -j ACCEPT 2>/dev/null || ip6tables -A INPUT -p tcp --dport "$port" -j ACCEPT
            ip6tables -C INPUT -p udp --dport "$port" -j ACCEPT 2>/dev/null || ip6tables -A INPUT -p udp --dport "$port" -j ACCEPT
        fi
        if [[ -e /etc/iptables/rules.v4 ]]; then
            iptables-save > /etc/iptables/rules.v4
            [[ -e /etc/iptables/rules.v6 ]] && ip6tables-save > /etc/iptables/rules.v6
        elif [[ -e /etc/sysconfig/iptables ]]; then
            iptables-save > /etc/sysconfig/iptables
            [[ -e /etc/sysconfig/ip6tables ]] && ip6tables-save > /etc/sysconfig/ip6tables
        fi
        echo "iptables: opened ${port}/tcp,udp"
        return 0
    fi

    echo "No active firewall detected, skipping."
}

function install_realm() {
    if [[ -f "$REALM_SERVICE" ]] || [[ -x "$REALM_BIN" ]]; then
        echo -e "${YELLOW}realm 已安装，如需重新安装请先卸载。${NC}"
        return 0
    fi
    check_dependencies
    prompt_install_dir || return 1
    prompt_proxy
    local version
    version=$(get_latest_version) || return 1
    download_realm "$version" || return 1
    write_default_config
    write_systemd_unit
    systemctl restart realm
    sleep 1
    if systemctl is-active --quiet realm; then
        echo -e "${GREEN}realm 安装完成并已启动。${NC}"
    else
        echo -e "${YELLOW}realm 已安装，但服务未启动 —— 当前没有任何转发规则。请通过菜单 [2] 添加规则。${NC}"
    fi
}

function ensure_installed() {
    if [[ ! -x "$REALM_BIN" ]]; then
        echo -e "${RED}realm 尚未安装，请先选择 [1] 安装。${NC}"
        return 1
    fi
    return 0
}

function list_rules() {
    ensure_installed || return 1
    if [[ ! -f "$REALM_CONF" ]]; then
        echo -e "${YELLOW}配置文件不存在。${NC}"
        return 0
    fi
    local idx=0
    local listen=""
    local remote=""
    echo "------------------------------------------------------------"
    printf " %-4s %-30s %-30s\n" "编号" "本地监听" "转发目标"
    echo "------------------------------------------------------------"
    while IFS= read -r line; do
        line="${line%%#*}"
        if [[ "$line" =~ ^[[:space:]]*\[\[endpoints\]\] ]]; then
            if [[ -n "$listen" || -n "$remote" ]]; then
                idx=$((idx+1))
                printf " %-4s %-30s %-30s\n" "$idx" "$listen" "$remote"
            fi
            listen=""
            remote=""
        elif [[ "$line" =~ ^[[:space:]]*listen[[:space:]]*=[[:space:]]*\"([^\"]+)\" ]]; then
            listen="${BASH_REMATCH[1]}"
        elif [[ "$line" =~ ^[[:space:]]*remote[[:space:]]*=[[:space:]]*\"([^\"]+)\" ]]; then
            remote="${BASH_REMATCH[1]}"
        fi
    done < "$REALM_CONF"
    if [[ -n "$listen" || -n "$remote" ]]; then
        idx=$((idx+1))
        printf " %-4s %-30s %-30s\n" "$idx" "$listen" "$remote"
    fi
    echo "------------------------------------------------------------"
    if [[ $idx -eq 0 ]]; then
        echo -e "${YELLOW}当前没有任何转发规则。${NC}"
    else
        echo "共 $idx 条规则。"
    fi
}

function add_rule() {
    ensure_installed || return 1
    write_default_config

    local listen_port listen_addr remote_host remote_port
    while true; do
        read -p "请输入本地监听端口 (1-65535): " listen_port
        if [[ "$listen_port" =~ ^[0-9]+$ ]] && (( listen_port >= 1 && listen_port <= 65535 )); then
            break
        fi
        echo -e "${RED}端口无效。${NC}"
    done

    read -p "请输入本地监听地址 [默认 0.0.0.0]: " listen_addr
    listen_addr="${listen_addr:-0.0.0.0}"

    while true; do
        read -p "请输入转发目标地址 (域名或 IP): " remote_host
        [[ -n "$remote_host" ]] && break
        echo -e "${RED}目标地址不能为空。${NC}"
    done

    while true; do
        read -p "请输入转发目标端口 (1-65535): " remote_port
        if [[ "$remote_port" =~ ^[0-9]+$ ]] && (( remote_port >= 1 && remote_port <= 65535 )); then
            break
        fi
        echo -e "${RED}端口无效。${NC}"
    done

    local listen="${listen_addr}:${listen_port}"
    local remote
    if [[ "$remote_host" =~ : && ! "$remote_host" =~ ^\[ ]]; then
        remote="[${remote_host}]:${remote_port}"
    else
        remote="${remote_host}:${remote_port}"
    fi

    cat >> "$REALM_CONF" <<EOF
[[endpoints]]
listen = "${listen}"
remote = "${remote}"

EOF

    open_firewall_port "$listen_port"
    systemctl restart realm
    sleep 1
    if systemctl is-active --quiet realm; then
        echo -e "${GREEN}规则已添加：${listen} -> ${remote}${NC}"
    else
        echo -e "${RED}规则已写入，但服务启动失败，请使用 [9] 查看日志排查。${NC}"
    fi
}

function delete_rule() {
    ensure_installed || return 1
    list_rules
    [[ ! -f "$REALM_CONF" ]] && return 0

    local total
    total=$(grep -cE '^\[\[endpoints\]\]' "$REALM_CONF")
    if (( total == 0 )); then
        return 0
    fi

    local idx
    read -p "请输入要删除的规则编号 (回车取消): " idx
    [[ -z "$idx" ]] && return 0
    if ! [[ "$idx" =~ ^[0-9]+$ ]] || (( idx < 1 || idx > total )); then
        echo -e "${RED}编号无效。${NC}"
        return 1
    fi

    awk -v target="$idx" '
        BEGIN { count = 0; skip = 0 }
        /^\[\[endpoints\]\]/ {
            count++
            if (count == target) { skip = 1; next }
            else { skip = 0 }
        }
        /^\[/ && !/^\[\[endpoints\]\]/ { skip = 0 }
        { if (!skip) print }
    ' "$REALM_CONF" > "${REALM_CONF}.new"

    awk 'BEGIN{blank=0} /^[[:space:]]*$/{blank++; if(blank<=1) print; next} {blank=0; print}' "${REALM_CONF}.new" > "${REALM_CONF}.tmp"
    mv "${REALM_CONF}.tmp" "$REALM_CONF"
    rm -f "${REALM_CONF}.new"

    systemctl restart realm
    echo -e "${GREEN}已删除规则 #${idx}${NC}"
}

function service_action() {
    ensure_installed || return 1
    local action="$1"
    systemctl "$action" realm
    sleep 1
    systemctl status realm --no-pager -l | head -n 15
}

function view_logs() {
    ensure_installed || return 1
    echo -e "${CYAN}按 Ctrl+C 退出日志查看。${NC}"
    journalctl -u realm -n 100 -f --no-pager
}

function update_realm() {
    ensure_installed || return 1
    local version current
    version=$(get_latest_version) || return 1
    current=$("$REALM_BIN" --version 2>/dev/null | awk '{print $NF}')
    echo "Current: ${current:-unknown} | Latest: ${version}"

    if [[ -n "$current" && "${version#v}" == "${current#v}" ]]; then
        read -p "已是最新版本，是否仍要重新下载? [y/N]: " confirm
        [[ "$confirm" != "y" && "$confirm" != "Y" ]] && return 0
    fi

    systemctl stop realm
    download_realm "$version" || { systemctl start realm; return 1; }
    systemctl start realm
    echo -e "${GREEN}realm 已更新到 ${version}${NC}"
}

function update_script() {
    local self="$0"
    local tmp
    tmp=$(mktemp)
    if curl -fsSL "$SCRIPT_RAW_URL" -o "$tmp"; then
        if [[ -s "$tmp" ]] && head -n1 "$tmp" | grep -q '^#!/bin/bash'; then
            install -m 0755 "$tmp" "$self"
            rm -f "$tmp"
            echo -e "${GREEN}脚本已更新，请重新运行。${NC}"
            exit 0
        fi
    fi
    rm -f "$tmp"
    echo -e "${RED}脚本更新失败。${NC}"
}

function uninstall_realm() {
    read -p "确认卸载 realm 并删除全部配置? [y/N]: " confirm
    [[ "$confirm" != "y" && "$confirm" != "Y" ]] && return 0

    systemctl stop realm 2>/dev/null
    systemctl disable realm 2>/dev/null
    rm -f "$REALM_SERVICE"
    systemctl daemon-reload
    rm -rf "$INSTALL_DIR"
    echo -e "${GREEN}realm 已卸载（已删除 ${INSTALL_DIR}）。${NC}"
}

function show_status() {
    if [[ ! -f "$REALM_SERVICE" ]] || [[ ! -x "$REALM_BIN" ]]; then
        echo -e "状态: ${YELLOW}未安装${NC}"
        return
    fi
    local ver
    ver=$("$REALM_BIN" --version 2>/dev/null | awk '{print $NF}')
    local active
    if systemctl is-active --quiet realm; then
        active="${GREEN}running${NC}"
    else
        active="${RED}stopped${NC}"
    fi
    local rules
    rules=$(grep -cE '^\[\[endpoints\]\]' "$REALM_CONF" 2>/dev/null || echo 0)
    echo -e "状态: 已安装 (${ver:-unknown}) | 目录: ${INSTALL_DIR} | 服务: ${active} | 规则数: ${rules}"
}

function main_menu() {
    clear
    echo "╔════════════════════════════════════════════════════════════════════════╗"
    echo -e "║ ${CYAN}realm 管理脚本${NC}                                                          ║"
    echo -e "║ ${CYAN}项目地址${NC}: https://github.com/zhboner/realm                              ║"
    echo "╠════════════════════════════════════════════════════════════════════════╣"
    show_status
    echo "╠════════════════════════════════════════════════════════════════════════╣"
    echo -e "║${CYAN} [1]${NC}  安装 realm                       ${CYAN} [2]${NC}  添加转发规则             ║"
    echo -e "║${CYAN} [3]${NC}  删除转发规则                     ${CYAN} [4]${NC}  查看转发规则             ║"
    echo -e "║${CYAN} [5]${NC}  启动服务                         ${CYAN} [6]${NC}  停止服务                 ║"
    echo -e "║${CYAN} [7]${NC}  重启服务                         ${CYAN} [8]${NC}  查看服务状态             ║"
    echo -e "║${CYAN} [9]${NC}  查看日志                         ${CYAN} [10]${NC} 更新 realm 内核          ║"
    echo -e "║${CYAN} [11]${NC} 更新脚本                         ${CYAN} [12]${NC} 卸载                     ║"
    echo -e "║${CYAN} [0]${NC}  退出                                                              ║"
    echo "╚════════════════════════════════════════════════════════════════════════╝"

    local choice
    read -p "请选择 [0-12]: " choice
    case "$choice" in
        1)  install_realm ;;
        2)  add_rule ;;
        3)  delete_rule ;;
        4)  list_rules ;;
        5)  service_action start ;;
        6)  service_action stop ;;
        7)  service_action restart ;;
        8)  ensure_installed && systemctl status realm --no-pager -l ;;
        9)  view_logs ;;
        10) update_realm ;;
        11) update_script ;;
        12) uninstall_realm ;;
        0)  exit 0 ;;
        *)  echo -e "${RED}无效的选择。${NC}" ;;
    esac
}

check_root
detect_install_dir
main_menu