toes/src/server/proxy.ts

import type { Server, ServerWebSocket } from 'bun'
import { getAppBySubdomain } from '$apps'
import { serveStatic } from '$static'

export const perf = { timing: false }

export type { WsData }

const pendingMessages = new Map<ServerWebSocket<WsData>, (string | ArrayBuffer | Uint8Array)[]>()
const upstreams = new Map<ServerWebSocket<WsData>, WebSocket>()

interface WsData {
  port: number
  path: string
  protocols: string[]
  headers: Record<string, string>
}

export function extractSubdomain(host: string): string | null {
  // Strip port
  const hostname = host.replace(/:\d+$/, '')

  // *.localhost -> take prefix (e.g. clock.localhost -> clock)
  if (hostname.endsWith('.localhost')) {
    const sub = hostname.slice(0, -'.localhost'.length)
    return sub || null
  }

  // *.X.local -> take first segment if 3+ parts (e.g. clock.toes.local -> clock)
  if (hostname.endsWith('.local')) {
    const parts = hostname.split('.')
    if (parts.length >= 3) {
      return parts[0]!
    }
  }

  return null
}

export async function proxySubdomain(subdomain: string, req: Request): Promise<Response> {
  const app = getAppBySubdomain(subdomain)

  if (!app || app.state !== 'running') {
    return new Response(`App "${subdomain}" not found or not running`, { status: 502 })
  }

  // Static apps: serve from pub/ directory
  if (app.static) {
    return serveStatic(app.name, req)
  }

  if (!app.port) {
    return new Response(`App "${subdomain}" not found or not running`, { status: 502 })
  }

  const url = new URL(req.url)
  const target = `http://localhost:${app.port}${url.pathname}${url.search}`
  const hasBody = ['POST', 'PUT', 'PATCH'].includes(req.method)
  const body = hasBody ? req.body : undefined

  const headers = new Headers(req.headers)
  headers.set('host', `localhost:${app.port}`)
  if (!headers.has('x-app-url')) {
    headers.set('x-app-url', app.tunnelUrl ?? `${url.protocol}//${subdomain}.${url.hostname}`)
  }
  headers.delete('connection')
  headers.delete('keep-alive')
  headers.delete('transfer-encoding')

  try {
    const shouldTime = perf.timing
    const start = shouldTime ? performance.now() : 0
    const res = await fetch(target, {
      method: req.method,
      headers,
      body,
      redirect: 'manual',
    })
    if (shouldTime) {
      const ms = (performance.now() - start).toFixed(1)
      console.log(`[perf] ${req.method} ${subdomain}${url.pathname} → ${res.status} in ${ms}ms`)
    }
    return res
  } catch (e) {
    console.error(`Proxy error for ${subdomain}:`, e)
    return new Response(`App "${subdomain}" is not responding`, { status: 502 })
  }
}

export function proxyWebSocket(subdomain: string, req: Request, server: Server<WsData>): Response | undefined {
  const app = getAppBySubdomain(subdomain)

  if (!app || app.state !== 'running' || !app.port) {
    return new Response(`App "${subdomain}" not found or not running`, { status: 502 })
  }

  const url = new URL(req.url)
  const path = url.pathname + url.search
  const protocolHeader = req.headers.get('sec-websocket-protocol')
  const protocols = protocolHeader ? protocolHeader.split(',').map(p => p.trim()) : []

  // Collect headers to forward to the upstream app
  const forwardHeaders: Record<string, string> = {}
  for (const name of ['cookie', 'authorization', 'x-app-url']) {
    const value = req.headers.get(name)
    if (value) forwardHeaders[name] = value
  }
  if (!forwardHeaders['x-app-url']) {
    forwardHeaders['x-app-url'] = app.tunnelUrl ?? `${url.protocol}//${subdomain}.${url.hostname}`
  }

  const upgradeHeaders: Record<string, string> = {}
  if (protocolHeader) upgradeHeaders['sec-websocket-protocol'] = protocolHeader

  const ok = server.upgrade(req, { data: { port: app.port, path, protocols, headers: forwardHeaders } as WsData, headers: upgradeHeaders })
  if (ok) return undefined
  return new Response('WebSocket upgrade failed', { status: 500 })
}

export const websocket = {
  open(ws: ServerWebSocket<WsData>) {
    const { port, path } = ws.data
    const upstream = new WebSocket(`ws://localhost:${port}${path}`, {
      headers: { ...ws.data.headers, host: `localhost:${port}` },
      protocols: ws.data.protocols,
    })

    upstream.binaryType = 'arraybuffer'
    upstreams.set(ws, upstream)
    pendingMessages.set(ws, [])

    const timeout = setTimeout(() => {
      if (upstream.readyState !== WebSocket.OPEN) {
        upstream.close()
        ws.close()
      }
    }, 10_000)

    upstream.addEventListener('open', () => {
      clearTimeout(timeout)
      const buffered = pendingMessages.get(ws)
      if (buffered) {
        for (const msg of buffered) upstream.send(msg)
        pendingMessages.delete(ws)
      }
    })

    upstream.addEventListener('message', e => {
      // binaryType is 'arraybuffer' so data is always string | ArrayBuffer
      ws.send(e.data as string | ArrayBuffer)
    })

    upstream.addEventListener('close', () => {
      clearTimeout(timeout)
      pendingMessages.delete(ws)
      upstreams.delete(ws)
      ws.close()
    })

    upstream.addEventListener('error', () => {
      clearTimeout(timeout)
      pendingMessages.delete(ws)
      upstreams.delete(ws)
      ws.close()
    })
  },

  message(ws: ServerWebSocket<WsData>, msg: string | ArrayBuffer | Uint8Array) {
    const upstream = upstreams.get(ws)
    if (!upstream) return
    if (upstream.readyState !== WebSocket.OPEN) {
      pendingMessages.get(ws)?.push(msg)
      return
    }
    upstream.send(msg)
  },

  close(ws: ServerWebSocket<WsData>) {
    const upstream = upstreams.get(ws)
    if (upstream) {
      upstream.close()
      upstreams.delete(ws)
    }
    pendingMessages.delete(ws)
  },
}