#!/bin/bash

# ================= 颜色定义 =================
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;36m'
PLAIN='\033[0m'

echo -e "${BLUE}==================================================${PLAIN}"
echo -e "${GREEN}  华为云 Token 动态网关 (6小时高能效缓存版) 部署  ${PLAIN}"
echo -e "${BLUE}==================================================${PLAIN}"

# 1. 检查权限
if [ "$EUID" -ne 0 ]; then
    echo -e "${RED}错误：请使用 root 用户运行此脚本！${PLAIN}"
    exit 1
fi

# ================= 检测并安装 frpc =================
echo -e "\n${YELLOW}[1/5] 正在检测 frpc 环境...${PLAIN}"
if ! command -v frpc &> /dev/null && [ ! -f "/usr/local/bin/frpc" ] && [ ! -f "/usr/bin/frpc" ]; then
    echo -e "${RED}未检测到 frpc！${PLAIN}"
    read -p "是否需要为您自动下载并安装最新版 frpc (使用国内加速镜像)? (y/n, 默认 y): " INSTALL_FRP
    INSTALL_FRP=${INSTALL_FRP:-y}
    
    if [[ "$INSTALL_FRP" == "y" || "$INSTALL_FRP" == "Y" ]]; then
        ARCH=$(uname -m)
        if [ "$ARCH" = "x86_64" ]; then
            FRP_ARCH="amd64"
        elif [ "$ARCH" = "aarch64" ]; then
            FRP_ARCH="arm64"
        else
            echo -e "${RED}错误：不支持的系统架构 $ARCH，请手动安装 frpc。${PLAIN}"
            exit 1
        fi
        
        FRP_VERSION="0.58.1"
        echo -e "${GREEN}开始通过国内节点下载 frpc v${FRP_VERSION} ($FRP_ARCH)...${PLAIN}"
        FRP_URL="https://mirror.ghproxy.com/https://github.com/fatedier/frp/releases/download/v${FRP_VERSION}/frp_${FRP_VERSION}_linux_${FRP_ARCH}.tar.gz"
        
        wget --no-check-certificate -O frp.tar.gz "$FRP_URL"
        if [ $? -ne 0 ]; then echo -e "${RED}下载失败！${PLAIN}"; exit 1; fi
        
        tar -zxvf frp.tar.gz > /dev/null
        cd frp_${FRP_VERSION}_linux_${FRP_ARCH} || { echo -e "${RED}错误：无法进入解压目录！${PLAIN}"; exit 1; }
        cp frpc /usr/local/bin/ && chmod +x /usr/local/bin/frpc
        mkdir -p /etc/frp
        
        if [ ! -f "/etc/frp/frpc.toml" ]; then
            cat << EOF > /etc/frp/frpc.toml
serverAddr = "x.x.x.x"
serverPort = 7000
EOF
        fi
        
        cat << EOF > /etc/systemd/system/frpc.service
[Unit]
Description=Frp Client Service
After=network.target

[Service]
Type=simple
User=root
Restart=on-failure
RestartSec=5s
ExecStart=/usr/local/bin/frpc -c /etc/frp/frpc.toml
ExecReload=/usr/local/bin/frpc reload -c /etc/frp/frpc.toml

[Install]
WantedBy=multi-user.target
EOF
        systemctl daemon-reload && systemctl enable frpc --now
        cd .. && rm -rf frp_${FRP_VERSION}_linux_${FRP_ARCH} frp.tar.gz
        echo -e "${GREEN}frpc 安装并注册成功！配置文件位于 /etc/frp/frpc.toml${PLAIN}"
    fi
else
    echo -e "${GREEN}系统已安装 frpc。${PLAIN}"
fi

# ================= 配置网关与映射端口 =================
echo -e "\n${YELLOW}[2/5] 配置动态网关参数：${PLAIN}"
read -p "请输入 Python 动态网关的本地监听端口 (默认 8080): " GATEWAY_PORT
GATEWAY_PORT=${GATEWAY_PORT:-8080}

read -p "请输入你想在 frps 公网端暴露的端口 (默认 18000): " REMOTE_PORT
REMOTE_PORT=${REMOTE_PORT:-18000}

# ================= 查找 frpc 配置文件 =================
echo -e "\n${YELLOW}[3/5] 检测 frpc 配置文件位置以追加代理规则：${PLAIN}"
FRPC_CONFIG=""
if [ -f "/etc/frp/frpc.toml" ]; then FRPC_CONFIG="/etc/frp/frpc.toml"; FRPC_TYPE="toml";
elif [ -f "/etc/frp/frpc.ini" ]; then FRPC_CONFIG="/etc/frp/frpc.ini"; FRPC_TYPE="ini";
else
    FRPC_PATH_SEARCH=$(find /etc /usr/local/etc /opt /home -maxdepth 4 -name "frpc.toml" -o -name "frpc.ini" 2>/dev/null | head -n 1)
    if [ -n "$FRPC_PATH_SEARCH" ]; then
        FRPC_CONFIG="$FRPC_PATH_SEARCH"
        if [[ "$FRPC_PATH_SEARCH" == *.toml ]]; then FRPC_TYPE="toml"; else FRPC_TYPE="ini"; fi
    fi
fi

if [ -n "$FRPC_CONFIG" ]; then
    echo -e "${GREEN}识别到 frpc 配置: $FRPC_CONFIG (${FRPC_TYPE})${PLAIN}"
    read -p "是否自动将华为云转发规则追加到该文件? (y/n, 默认 y): " APPEND_FRP
    APPEND_FRP=${APPEND_FRP:-y}
else
    APPEND_FRP="n"
fi

# ================= 安装 Python 依赖 =================
echo -e "\n${YELLOW}[4/5] 部署 Python 运行环境...${PLAIN}"
if ! command -v python3 &> /dev/null; then
    if command -v apt &> /dev/null; then apt update -y && apt install -y python3 python3-pip;
    elif command -v yum &> /dev/null; then yum install -y python3 python3-pip; fi
fi
python3 -m pip install flask requests --quiet -i https://pypi.tuna.tsinghua.edu.cn/simple

# ================= 创建网关核心代码 (核心优化点) =================
echo -e "\n${YELLOW}[5/5] 生成网关代码并启动服务...${PLAIN}"
cat << 'EOF' > /usr/local/bin/huawei_gateway.py
import os
import re
import time
import logging
import threading
import traceback
import requests
from flask import Flask, request, Response
from concurrent.futures import ThreadPoolExecutor

app = Flask(__name__)
logging.basicConfig(level=logging.WARNING,
                    format='%(asctime)s [%(levelname)s] %(message)s')

# 内存缓存字典：6小时有效期，这里设置为安全线 5.5 小时 (19800秒)
cache = {"token": None, "expires_at": 0}
cache_lock = threading.Lock()

def scan_pid_mem(pid):
    try:
        with open(f'/proc/{pid}/maps', 'r') as f:
            for line in f:
                parts = line.split()
                if len(parts) < 2 or 'r' not in parts[1]: continue
                addrs = parts[0].split('-')
                if len(addrs) != 2: continue
                start, end = int(addrs[0], 16), int(addrs[1], 16)
                if (end - start) > 200 * 1024 * 1024: continue
                try:
                    with open(f'/proc/{pid}/mem', 'rb') as mem:
                        mem.seek(start)
                        data = mem.read(end - start)
                        for m in re.finditer(b'Bearer ([A-Za-z0-9+/=_-]+)', data):
                            val = m.group(1).decode(errors='replace')
                            if len(val) > 200:
                                return val
                except Exception:
                    continue
    except Exception:
        pass
    return None

def force_refresh_token():
    """强制重新从内存扫描 Token 并更新缓存时间"""
    pids = [pid for pid in os.listdir('/proc') if pid.isdigit()]
    with ThreadPoolExecutor(max_workers=8) as executor:
        results = executor.map(scan_pid_mem, pids)
        for token in results:
            if token:
                with cache_lock:
                    cache["token"] = token
                    cache["expires_at"] = time.time() + 19800  # 缓存 5.5 小时
                return token
    return cache["token"]

def get_latest_token():
    now = time.time()
    with cache_lock:
        # 如果缓存存在且没过期，直接返回缓存（消耗几乎为 0）
        if cache["token"] and now < cache["expires_at"]:
            return cache["token"]
    # 否则，或者过期了，才去扫描内存
    return force_refresh_token()

@app.route('/v2/<path:subpath>', methods=['POST', 'GET', 'OPTIONS'])
def proxy(subpath):
    if request.method == 'OPTIONS':
        return Response(), 200

    real_token = get_latest_token()
    if not real_token:
        logging.error("未在内存中找到华为云Token")
        return {"error": "未在内存中找到华为云Token"}, 500

    headers = {k: v for k, v in request.headers if k.lower() != 'host'}
    headers['Authorization'] = f'Bearer {real_token}'
    headers['Host'] = 'tokenhub.developer.huaweicloud.com'

    target_url = f'https://tokenhub.developer.huaweicloud.com/v2/{subpath}'
    try:
        resp = requests.request(
            method=request.method, url=target_url, headers=headers,
            data=request.get_data(), cookies=request.cookies,
            allow_redirects=False, timeout=60
        )

        # 兜底保障：如果华为主端提早失效返回了 401，立即强制重新扫描内存并重试一次请求
        if resp.status_code == 401:
            new_token = force_refresh_token()
            if new_token:
                headers['Authorization'] = f'Bearer {new_token}'
                resp = requests.request(
                    method=request.method, url=target_url, headers=headers,
                    data=request.get_data(), cookies=request.cookies,
                    allow_redirects=False, timeout=60
                )

        return Response(resp.content, resp.status_code, resp.headers.items())
    except Exception as e:
        logging.error(f"网关转发失败: {traceback.format_exc()}")
        return {"error": f"网关转发失败: {str(e)}"}, 500

if __name__ == '__main__':
    import sys
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 8080
    app.run(host='127.0.0.1', port=port, debug=False)

EOF

# 注册网关服务
cat << EOF > /etc/systemd/system/huawei-gateway.service
[Unit]
Description=Huawei Dynamic Gateway (6H Cache)
After=network.target

[Service]
Type=simple
User=root
ExecStart=/usr/bin/python3 /usr/local/bin/huawei_gateway.py $GATEWAY_PORT
Restart=always
RestartSec=3

[Install]
WantedBy=multi-user.target
EOF

systemctl daemon-reload && systemctl enable huawei-gateway --now &>/dev/null

# ================= 追加 frpc 规则 =================
if [[ "$APPEND_FRP" == "y" || "$APPEND_FRP" == "Y" ]] && [ -f "$FRPC_CONFIG" ]; then
    if grep -q "huawei-dynamic-gateway" "$FRPC_CONFIG" 2>/dev/null; then
        echo -e "${YELLOW}检测到 frpc 配置中已存在 huawei-dynamic-gateway 规则，跳过追加。${PLAIN}"
    else
        if [ "$FRPC_TYPE" = "toml" ]; then
            cat << EOF >> "$FRPC_CONFIG"

[[proxies]]
name = "huawei-dynamic-gateway"
type = "tcp"
local_ip = "127.0.0.1"
local_port = $GATEWAY_PORT
remote_port = $REMOTE_PORT
EOF
        else
            cat << EOF >> "$FRPC_CONFIG"

[huawei-dynamic-gateway]
type = tcp
local_ip = 127.0.0.1
local_port = $GATEWAY_PORT
remote_port = $REMOTE_PORT
EOF
        fi
        echo -e "${GREEN}frpc 代理规则已追加到 $FRPC_CONFIG${PLAIN}"
    fi
    if systemctl list-units --type=service | grep -q "frpc"; then systemctl restart frpc; fi
fi

# ================= 部署总结 =================
echo -e "\n${BLUE}==================================================${PLAIN}"
echo -e "${GREEN}          🎉 6小时长效缓存网关部署完成 🎉         ${PLAIN}"
echo -e "${BLUE}==================================================${PLAIN}"
echo -e "1. 动态网关状态：$(systemctl is-active huawei-gateway)"
echo -e "2. 您的本地一劳永逸调用地址："
echo -e "   - ${GREEN}API Base URL:${PLAIN} http://【你的公网frps_IP】:$REMOTE_PORT/v2"
echo -e "   - ${GREEN}API Key:${PLAIN} sk-anything"
echo -e "   - ${GREEN}Model:${PLAIN} glm-5.1"
echo -e "${BLUE}==================================================${PLAIN}\n"