toes/src/server/proxy.ts

import type { Server, ServerWebSocket } from 'bun'
import { getApp } from '$apps'

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[]
}

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 = getApp(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 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}`)
  headers.delete('connection')
  headers.delete('content-length')
  headers.delete('keep-alive')
  headers.delete('transfer-encoding')

  try {
    return await fetch(target, {
      method: req.method,
      headers,
      body,
      redirect: 'manual',
    })
  } 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 = getApp(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()) : []

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

  const ok = server.upgrade(req, { data: { port: app.port, path, protocols } as WsData, headers })
  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}`, 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)
  },
}