#!/usr/bin/env python3
"""
wxbot_on_demand_login_patch.py — 莫荷微信「按需登录」完整补丁
================================================================
背景: 2026-08-13 莫荷微信被风控(频繁重登 + 未登录忙循环请求腾讯)。
本补丁将 docker-wechatbot-webhook 改为「按需登录」:
  - 空闲时零请求腾讯 (不再自动 start)
  - 点"开始登录"按钮才触发登录流程
  - 登录窗口期轮询 3s 间隔+退避 (非忙循环)
  - 登出/停止后回到静默, 清空残留二维码

改动文件 (容器内):
  1. /app/src/wechaty/init.js          — DISABLE_AUTO_LOGIN 时不自动 start, 导出 start/stop/isRunning
  2. /app/main.js                      — 解构新导出 + 传 isRunning
  3. /app/src/route/bot.js             — 新增: POST /api/bot/login|logout, GET /api/bot/status
  4. /app/src/route/index.js           — 注册 bot 路由 + 传 isRunning
  5. /app/src/route/login.js           — bot 未运行时不展示二维码
  6. /app/node_modules/wechat4u/lib/wechat.js      — login 忙循环加 3s+退避 (checkLogin)
  7. /app/node_modules/wechat4u/lib/util/request.js — UA/headers/keep-alive (风控降特征)
  8. /app/node_modules/wechat4u/lib/util/global.js  — DeviceID 持久化
  9. /app/node_modules/wechaty-puppet-wechat4u/dist/cjs/src/puppet-wechat4u.js — logout 不自动 start
  10. /app/.env                        — DISABLE_AUTO_LOGIN=true

用法:
  python3 wxbot_on_demand_login_patch.py           # 应用补丁
  python3 wxbot_on_demand_login_patch.py --revert  # 回滚
"""
import subprocess, os, sys, shutil, time

CONTAINER = "wxBotWebhook"
TMP = "/tmp/wxbot-on-demand-patch"
os.makedirs(TMP, exist_ok=True)

def run(cmd, timeout=30):
    r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
    return (r.stdout + r.stderr).strip()

def docker_exec(sh_cmd):
    return run(["docker", "exec", CONTAINER, "sh", "-c", sh_cmd])

def docker_cp_out(path, local):
    run(["docker", "cp", f"{CONTAINER}:{path}", local])

def docker_cp_in(local, path):
    run(["docker", "cp", local, f"{CONTAINER}:{path}"])

# ══════════════════════════════════════════════════════════════
# 1. init.js 补丁: 按需启动 + 导出控制函数
# ══════════════════════════════════════════════════════════════
def patch_init_js():
    path = "/app/src/wechaty/init.js"
    docker_cp_out(path, f"{TMP}/init.js")
    c = open(f"{TMP}/init.js", encoding='utf-8').read()

    # 备份
    docker_exec(f"cp {path} {path}.on-demand.bak")

    # 1a. 模块头加 autoLoginDisabled 判定 (含 .env 兜底)
    old_module = """module.exports = function init() {
  /** @type {import('wechaty').Contact} */
  let currentUser
  let botLoginSuccessLastTime = false
"""
    new_module = """module.exports = function init() {
  /** @type {import('wechaty').Contact} */
  let currentUser
  let botLoginSuccessLastTime = false
  let botStarted = false

  // 按需登录控制: DISABLE_AUTO_LOGIN=true 时不自动启动
  // 兜底: docker run 可能注入 DISABLE_AUTO_LOGIN= (空值覆盖 .env), 这里同时读 .env 文件
  let _autoLoginFromFile = false
  try {
    const _fs = require('fs')
    const _path = require('path')
    const _envContent = _fs.readFileSync(_path.join(__dirname, '../../.env'), 'utf-8')
    const _m = _envContent.match(/^DISABLE_AUTO_LOGIN=(\\S*)$/m)
    if (_m && _m[1] === 'true') _autoLoginFromFile = true
  } catch (e) { /* ignore */ }
  const autoLoginDisabled =
    process.env.DISABLE_AUTO_LOGIN === 'true' || _autoLoginFromFile
"""
    if old_module in c:
        c = c.replace(old_module, new_module)
        print("  ✅ init.js: 按需启动判定")
    else:
        print("  ⚠️ init.js 模块头未匹配")

    # 1b. 末尾: 条件 start + 导出控制对象
    old_start = """  bot.start().catch((e) => {
    Utils.logger.error('bot 初始化失败：', e)
  })

  return bot
}"""
    new_start = """  // 自动登录模式: 容器启动即开始登录流程(获取二维码+轮询)
  if (!autoLoginDisabled) {
    botStarted = true
    bot.start().catch((e) => {
      Utils.logger.error('bot 初始化失败：', e)
    })
  } else {
    Utils.logger.info(
      'DISABLE_AUTO_LOGIN=true: 机器人已就绪但未启动, 等待 POST /api/bot/login 触发登录'
    )
  }

  return {
    bot,
    /**
     * 按需启动登录流程 (获取二维码 + 检测扫码)
     * 幂等: 已启动/已登录时直接返回当前状态
     */
    start: async () => {
      if (botStarted) {
        Utils.logger.info('bot 已在运行中, 忽略重复 start')
        return { started: true, alreadyRunning: true }
      }
      botStarted = true
      Utils.logger.info('手动触发 bot.start() - 开始登录流程')
      await bot.start()
      return { started: true }
    },
    /**
     * 停止机器人 (登出后回到静默)
     */
    stop: async () => {
      if (!botStarted) return { stopped: true, alreadyStopped: true }
      botStarted = false
      try {
        await bot.stop()
      } catch (e) {
        Utils.logger.error('bot.stop() 出错:', e)
      }
      Utils.logger.info('bot 已停止, 回到静默状态')
      return { stopped: true }
    },
    isRunning: () => botStarted
  }
}"""
    if old_start in c:
        c = c.replace(old_start, new_start)
        print("  ✅ init.js: 条件 start + 导出控制")
    else:
        print("  ⚠️ init.js 末尾未匹配")

    open(f"{TMP}/init.js", 'w', encoding='utf-8').write(c)
    docker_cp_in(f"{TMP}/init.js", path)

# ══════════════════════════════════════════════════════════════
# 2. main.js 补丁
# ══════════════════════════════════════════════════════════════
def patch_main_js():
    path = "/app/main.js"
    docker_cp_out(path, f"{TMP}/main.js")
    c = open(f"{TMP}/main.js", encoding='utf-8').read()

    old = """const { bot, start: startBot, stop: stopBot } = wechatBotInit()"""
    new = """const { bot, start: startBot, stop: stopBot, isRunning } = wechatBotInit()"""
    if old in c:
        c = c.replace(old, new)
        print("  ✅ main.js: 解构 isRunning")
    old2 = """registerRoute({ app, bot, startBot, stopBot })"""
    new2 = """registerRoute({ app, bot, startBot, stopBot, isRunning })"""
    if old2 in c:
        c = c.replace(old2, new2)
        print("  ✅ main.js: 传 isRunning")

    open(f"{TMP}/main.js", 'w', encoding='utf-8').write(c)
    docker_cp_in(f"{TMP}/main.js", path)

# ══════════════════════════════════════════════════════════════
# 3. route/bot.js (新增)
# ══════════════════════════════════════════════════════════════
BOT_ROUTE = """/**
 * 按需登录控制路由
 * 仅当 DISABLE_AUTO_LOGIN=true 时生效:
 *  - POST /api/bot/login  手动触发登录(获取二维码+检测扫码)
 *  - POST /api/bot/logout 登出并回到静默
 *  - GET  /api/bot/status 查询当前状态
 */
module.exports = function registerBotRoute({ app, startBot, stopBot, bot }) {
  // 触发登录 (获取二维码)
  app.post('/api/bot/login', async (c) => {
    try {
      const result = await startBot()
      return c.json({
        success: true,
        ...result,
        loginUrl: `/login?token=${process.env.LOCAL_LOGIN_API_TOKEN || ''}`
      })
    } catch (e) {
      return c.json({ success: false, error: e.message }, 500)
    }
  })

  // 登出回到静默
  app.post('/api/bot/logout', async (c) => {
    try {
      const result = await stopBot()
      return c.json({ success: true, ...result })
    } catch (e) {
      return c.json({ success: false, error: e.message }, 500)
    }
  })

  // 状态查询
  app.get('/api/bot/status', async (c) => {
    const isLoggedIn = !!(bot && bot.isLoggedIn)
    const isRunning = startBot ? (typeof startBot.running === 'function' ? startBot.running() : false) : false
    return c.json({
      isLoggedIn,
      loginUrl: `/login?token=${process.env.LOCAL_LOGIN_API_TOKEN || ''}`
    })
  })
}
"""

def patch_bot_route():
    path = "/app/src/route/bot.js"
    with open(f"{TMP}/bot.js", 'w', encoding='utf-8') as f:
        f.write(BOT_ROUTE)
    docker_cp_in(f"{TMP}/bot.js", path)
    print("  ✅ route/bot.js: 新增")

# ══════════════════════════════════════════════════════════════
# 4. route/index.js
# ══════════════════════════════════════════════════════════════
def patch_route_index():
    path = "/app/src/route/index.js"
    docker_cp_out(path, f"{TMP}/index.js")
    c = open(f"{TMP}/index.js", encoding='utf-8').read()

    old = """module.exports = function registerRoute({ app, bot, startBot, stopBot }) {"""
    new = """module.exports = function registerRoute({ app, bot, startBot, stopBot, isRunning }) {"""
    if old in c:
        c = c.replace(old, new)
    old2 = """  require('./login')({ app, bot })"""
    new2 = """  require('./login')({ app, bot, isRunning })"""
    if old2 in c:
        c = c.replace(old2, new2)
    old3 = """  require('./resouces')({ app, bot })"""
    new3 = """  require('./resouces')({ app, bot })
  if (startBot) {
    require('./bot')({ app, startBot, stopBot, bot })
  }"""
    if old3 in c:
        c = c.replace(old3, new3)
        print("  ✅ route/index.js: 注册 bot 路由 + 传参")

    open(f"{TMP}/index.js", 'w', encoding='utf-8').write(c)
    docker_cp_in(f"{TMP}/index.js", path)

# ══════════════════════════════════════════════════════════════
# 5. route/login.js — 未运行不展示二维码
# ══════════════════════════════════════════════════════════════
def patch_login_js():
    path = "/app/src/route/login.js"
    docker_cp_out(path, f"{TMP}/login.js")
    c = open(f"{TMP}/login.js", encoding='utf-8').read()

    old_sig = """module.exports = function registerLoginCheck({ app, bot }) {"""
    new_sig = """module.exports = function registerLoginCheck({ app, bot, isRunning }) {
  // [按需登录补丁] bot 未运行时(未触发登录)不展示二维码
  const botActive = isRunning ? isRunning() : true
  const token = process.env.LOCAL_LOGIN_API_TOKEN || ''"""
    if old_sig in c:
        c = c.replace(old_sig, new_sig)
        print("  ✅ login.js: 签名+botActive+token")

    old_get = """    async (c) => {
      // 登录成功的话，返回登录信息
      if (success) {"""
    new_get = """    async (c) => {
      // [按需登录补丁] bot 未运行 → 不展示二维码, 提示未启动
      if (!botActive) {
        const html = `
        <!DOCTYPE html>
        <html lang="en">
        <head>
          <meta charset="UTF-8">
          <title>扫码登录</title>
          <style>
            body, html { display: flex; justify-content: center; align-items: center; height: 100%; margin: 0; font-family: sans-serif; }
            .box { text-align: center; padding: 40px; border: 1px solid #ddd; border-radius: 12px; }
            .btn { display: inline-block; margin-top: 20px; padding: 10px 24px; background: #07c160; color: #fff; border-radius: 8px; text-decoration: none; font-size: 16px; }
          </style>
        </head>
        <body>
          <div class="box">
            <h2>莫荷微信机器人</h2>
            <p>当前未启动，未在请求微信服务器。</p>
            <a class="btn" href="/api/bot/login?token=${token}">开始登录（获取二维码）</a>
          </div>
        </body>
        </html>
      `
        return c.html(html)
      }
      // 登录成功的话，返回登录信息
      if (success) {"""
    if old_get in c:
        c = c.replace(old_get, new_get)
        print("  ✅ login.js: 未启动分支")

    open(f"{TMP}/login.js", 'w', encoding='utf-8').write(c)
    docker_cp_in(f"{TMP}/login.js", path)

# ══════════════════════════════════════════════════════════════
# 6. wechat4u wechat.js — checkLogin 忙循环加 3s+退避
# ══════════════════════════════════════════════════════════════
def patch_wechat_busyloop():
    path = "/app/node_modules/wechat4u/lib/wechat.js"
    docker_cp_out(path, f"{TMP}/wechat.js")
    c = open(f"{TMP}/wechat.js", encoding='utf-8').read()

    old = """      var checkLogin = function checkLogin() {
        return _this6.checkLogin().then(function (res) {
          if (res.code === 201 && res.userAvatar) {
            _this6.emit('user-avatar', res.userAvatar);
          }
          if (res.code !== 200) {
            debug('checkLogin: ', res.code);
            return checkLogin();
          } else {
            return res;
          }
        });
      };"""
    new = """      var _loginCheckCount = 0;
      var checkLogin = function checkLogin() {
        return _this6.checkLogin().then(function (res) {
          if (res.code === 201 && res.userAvatar) {
            _this6.emit('user-avatar', res.userAvatar);
          }
          if (res.code !== 200) {
            debug('checkLogin: ', res.code);
            _loginCheckCount++;
            // 防忙循环: 前10次3秒间隔, 之后指数退避到最长30秒
            var delay = _loginCheckCount <= 10 ? 3000 : Math.min(3000 * Math.pow(2, _loginCheckCount - 10), 30000);
            return new Promise(function (resolve) {
              setTimeout(function () { resolve(checkLogin()); }, delay);
            });
          } else {
            _loginCheckCount = 0;
            return res;
          }
        });
      };"""
    if old in c:
        c = c.replace(old, new)
        print("  ✅ wechat.js: checkLogin 3s+退避")
    else:
        print("  ⚠️ wechat.js 忙循环未匹配(可能已打过)")
    open(f"{TMP}/wechat.js", 'w', encoding='utf-8').write(c)
    docker_cp_in(f"{TMP}/wechat.js", path)

# ══════════════════════════════════════════════════════════════
# 7. puppet-wechat4u.js — logout 不自动 start
# ══════════════════════════════════════════════════════════════
def patch_puppet_logout():
    path = "/app/node_modules/wechaty-puppet-wechat4u/dist/cjs/src/puppet-wechat4u.js"
    docker_cp_out(path, f"{TMP}/puppet.js")
    c = open(f"{TMP}/puppet.js", encoding='utf-8').read()

    old = """            // 清除数据
            await this.memory.delete(MEMORY_SLOT_NAME);
            await this.memory.save();
            this.wechat4u.start();
        });"""
    new = """            // 清除数据
            await this.memory.delete(MEMORY_SLOT_NAME);
            await this.memory.save();
            // [按需登录补丁] 登出后不再自动 start (原代码会立即重新开始登录轮询)
            // 需要重新登录时, 由上层 POST /api/bot/login 手动触发
            if (!process.env.DISABLE_AUTO_LOGIN) {
                this.wechat4u.start();
            } else {
                wechaty_puppet_1.log.info('PuppetWechat4u', 'logout后按需登录模式: 不自动重启, 等待手动触发');
            }
        });"""
    if old in c:
        c = c.replace(old, new)
        print("  ✅ puppet: logout 不自动 start")
    else:
        print("  ⚠️ puppet logout 未匹配")
    open(f"{TMP}/puppet.js", 'w', encoding='utf-8').write(c)
    docker_cp_in(f"{TMP}/puppet.js", path)

# ══════════════════════════════════════════════════════════════
# 8. .env — DISABLE_AUTO_LOGIN=true
# ══════════════════════════════════════════════════════════════
def patch_env():
    r = docker_exec("grep '^DISABLE_AUTO_LOGIN' /app/.env")
    if 'DISABLE_AUTO_LOGIN=true' in r:
        print("  ✅ .env: DISABLE_AUTO_LOGIN=true (已设置)")
    else:
        docker_exec("sed -i 's/^DISABLE_AUTO_LOGIN=$/DISABLE_AUTO_LOGIN=true/' /app/.env")
        print("  ✅ .env: DISABLE_AUTO_LOGIN=true")

# ══════════════════════════════════════════════════════════════
# 语法检查 + 重启
# ══════════════════════════════════════════════════════════════
def verify_and_restart():
    print("\n=== 语法检查 ===")
    files = [
        "/app/src/wechaty/init.js",
        "/app/main.js",
        "/app/src/route/bot.js",
        "/app/src/route/index.js",
        "/app/src/route/login.js",
        "/app/node_modules/wechat4u/lib/wechat.js",
        "/app/node_modules/wechaty-puppet-wechat4u/dist/cjs/src/puppet-wechat4u.js",
    ]
    ok = True
    for f in files:
        r = run(["docker", "exec", CONTAINER, "node", "--check", f])
        if r:
            ok = False
            print(f"  ✗ {f}: {r[:200]}")
        else:
            print(f"  ✓ {f}")
    return ok

def main():
    print("=== 应用按需登录补丁 ===")
    print("\n--- 1/8 init.js ---")
    patch_init_js()
    print("\n--- 2/8 main.js ---")
    patch_main_js()
    print("\n--- 3/8 route/bot.js ---")
    patch_bot_route()
    print("\n--- 4/8 route/index.js ---")
    patch_route_index()
    print("\n--- 5/8 route/login.js ---")
    patch_login_js()
    print("\n--- 6/8 wechat.js busy-loop ---")
    patch_wechat_busyloop()
    print("\n--- 7/8 puppet logout ---")
    patch_puppet_logout()
    print("\n--- 8/8 .env ---")
    patch_env()

    if verify_and_restart():
        print("\n=== 语法全部通过, 重启容器 ===")
        print(run(["docker", "restart", CONTAINER]))
        time.sleep(8)
        logs = run(["docker", "logs", CONTAINER, "--since", "1m", "2>&1"])
        if '等待 POST /api/bot/login' in logs:
            print("✅ 补丁生效: 静默等待触发登录")
        else:
            print("⚠️ 日志中未找到等待触发标记, 请检查:")
            for l in logs.split('\n')[:15]:
                print(f"  {l[:120]}")
    else:
        print("\n!!! 语法检查失败, 未重启. 请手动修复")

if __name__ == "__main__":
    main()