#!/usr/bin/env python3
# 档案 24 补丁：实例侧 401（launch token 过期）在浏览器导航时自动恢复
# 用法: python3 patch-24.py /opt/dshs
import io
import os
import sys

root = sys.argv[1] if len(sys.argv) > 1 else '.'
p = os.path.join(root, 'src/supervisor/proxy.ts')
s = io.open(p, encoding='utf-8').read()
changed = []


def sub(old, new, cnt=1):
    global s
    assert old in s, 'anchor not found:\n' + old[:200]
    s = s.replace(old, new, cnt)
    changed.append(old.split('\n')[0][:60])


# 1) 成功分支带上 userId（用于取新 token）
sub(
    """  const endpoint = await app.supervisor.endpointFor(session.user.id)
  if (endpoint === undefined) return { error: 'not_running', code: 404, userId: session.user.id }""",
    """  const endpoint = await app.supervisor.endpointFor(session.user.id)
  if (endpoint === undefined) return { error: 'not_running', code: 404, userId: session.user.id }
  // userId 一并返回：实例侧 401（launch token 过期）时用它取当前实例的新 token（档案 24）。""",
)

# 2) proxyHttp 增加 freshAuthUrl 参数
sub(
    """  rewriteLoopbackLocation = false,
  useKeepAlive = false,
): void {
  reply.hijack()""",
    """  rewriteLoopbackLocation = false,
  useKeepAlive = false,
  /**
   * 实例侧 401 时的恢复入口（档案 24）：浏览器导航遇到 dsh 的 "authentication required"
   * （launch token 过期 —— 实例重启/回收后刷新旧标签页）时调用，返回应 302 到的地址。
   * 返回 undefined 则回落到 '/'。
   */
  freshAuthUrl?: () => Promise<string | undefined>,
): void {
  reply.hijack()""",
)

# 3) 上游响应里的 401 处理（插在 writeHead 之前）
sub(
    """        reply.raw.writeHead(upRes.statusCode ?? 502, headers)""",
    """        // 2026-09-11（档案 24）：实例侧 401 = launch token 过期。浏览器导航时不要停在
        // dsh 的 "dsh web authentication required; reopen the URL printed by dsh web." 死端页，
        // 而是用当前实例的新 token 302 回同一地址（拿不到则回门户）。
        if (
          upRes.statusCode === 401 &&
          request.raw.method === 'GET' &&
          String(request.headers.accept ?? '').includes('text/html') &&
          freshAuthUrl !== undefined
        ) {
          upRes.resume()
          void freshAuthUrl()
            .then((url) => {
              reply.raw.writeHead(302, { location: url ?? '/' })
              reply.raw.end()
            })
            .catch(() => {
              reply.raw.writeHead(302, { location: '/' })
              reply.raw.end()
            })
          return
        }
        reply.raw.writeHead(upRes.statusCode ?? 502, headers)""",
)

# 4) 子域 onRequest 调用点：传入 freshAuthUrl
sub(
    """    proxyHttp(request, reply, access.endpoint, request.raw.url ?? '/', undefined, app.config.deployMode === 'local', app.config.deployMode === 'k8s')
  })

  // Per-user subdomain: WebSocket upgrade tunnel.""",
    """    const navScheme = app.config.secureCookies ? 'https' : 'http'
    const navHost = clientHost(request.headers) ?? app.config.baseDomain
    proxyHttp(
      request,
      reply,
      access.endpoint,
      request.raw.url ?? '/',
      undefined,
      app.config.deployMode === 'local',
      app.config.deployMode === 'k8s',
      async () => {
        const status = await app.supervisor.status(access.userId)
        const token = status.main?.launchToken
        return token !== undefined && token !== ''
          ? `${navScheme}://${navHost}/?token=${encodeURIComponent(token)}`
          : `${navScheme}://${app.config.baseDomain}/`
      },
    )
  })

  // Per-user subdomain: WebSocket upgrade tunnel.""",
)

# 5) 路径式路由 /u/:slug/dsh/ 调用点：同样传入
sub(
    """    proxyHttp(request, reply, endpoint, targetPath, prefix, app.config.deployMode === 'local', app.config.deployMode === 'k8s')
  })""",
    """    const pathScheme = app.config.secureCookies ? 'https' : 'http'
    proxyHttp(
      request,
      reply,
      endpoint,
      targetPath,
      prefix,
      app.config.deployMode === 'local',
      app.config.deployMode === 'k8s',
      async () => {
        const status = await app.supervisor.status(request.user!.id)
        const token = status.main?.launchToken
        return token !== undefined && token !== ''
          ? `${pathScheme}://${app.config.baseDomain}${prefix}/?token=${encodeURIComponent(token)}`
          : `${pathScheme}://${app.config.baseDomain}/`
      },
    )
  })""",
)

io.open(p, 'w', encoding='utf-8', newline='').write(s)
print('patched anchors:')
for c in changed:
    print('  -', c)