brf/server/plugins/vite/production.ts

import path from 'node:path'
import fs from 'node:fs'
import { type Readable } from 'node:stream'
import { resolveConfig, type ResolvedConfig as ViteResolvedConfig, type Manifest as ViteManifest } from 'vite'
import type { FastifyInstance, RouteHandler } from 'fastify'
import fstatic from '@fastify/static'

import StatusError from '../../lib/status_error.ts'

import type { Route } from '../../../shared/types.ts'
import type { Entry, ParsedConfig, Renderer, RenderCache, RenderFunction, ViteEntryFile } from './types.ts'

export default async function viteProduction(fastify: FastifyInstance, config: ParsedConfig) {
  const viteConfig = await resolveConfig({}, 'build', 'production')

  fastify.register(fstatic, {
    root: viteConfig.build.outDir,
    wildcard: false,
  })

  fastify.decorateReply('ctx', null)

  for (const entry of Object.values(config.entries)) {
    fastify.register((fastify, _, done) => setupEntry(fastify, entry, config, viteConfig).then(() => done()), {
      prefix: entry.path,
    })
  }
}

async function setupEntry(
  fastify: FastifyInstance,
  entry: Entry,
  _config: ParsedConfig,
  viteConfig: ViteResolvedConfig,
) {
  const viteEntryFilePath = path.join(viteConfig.build.outDir, entry.name + '.js')
  const viteEntryFile: ViteEntryFile = fs.existsSync(viteEntryFilePath)
    ? await import(viteEntryFilePath)
    : await import(path.join(viteConfig.root, entry.name, 'server.ts'))
  const manifestPath = path.join(viteConfig.build.outDir, `manifest.${entry.name}.json`)
  const manifest: ViteManifest = fs.existsSync(manifestPath)
    ? (await import(manifestPath, { with: { type: 'json' } })).default
    : null
  const routes = entry.routes || viteEntryFile.routes

  if (!routes) throw new Error('No routes found')

  const cache = new Map()

  const renderer = createRenderer(fastify, entry, viteEntryFile.render, manifest)
  const cachedHandler = createCachedHandler(cache)
  const handler = createHandler(renderer)

  if (entry.preHandler || entry.createPreHandler) {
    fastify.addHook('onRequest', (_request, reply, done) => {
      reply.ctx = Object.create(null)

      done()
    })
  }

  if (entry.preHandler) {
    fastify.addHook('preHandler', entry.preHandler)
  }

  for (const route of routes.flatMap((route) => route.routes || route)) {
    const preHandler = entry.createPreHandler?.(route, entry)

    if (route.locales?.length) {
      const locales = [{ hostname: entry.hostname, path: route.path }, ...route.locales]

      for (const locale of locales) {
        fastify.route({
          method: 'GET',
          url: locale.path,
          onRequest(request, reply, done) {
            if (request.hostname !== locale.hostname) {
              // TODO should probably redirect to correct path on request.hostname, not to locale.hostname
              return reply.redirect('http://' + locale.hostname + request.url)
            }

            done()
          },
          handler: route.cache ? cachedHandler : handler,
          preHandler,
        })
      }
    } else {
      fastify.route({
        method: 'GET',
        url: route.path,
        handler: route.cache ? cachedHandler : handler,
        preHandler,
      })
    }
  }

  entry.setupRebuildCache?.(() => buildCache(entry, routes, renderer, cache))

  fastify.setNotFoundHandler(() => {
    throw new StatusError(404)
  })

  if (entry.createErrorHandler) {
    fastify.setErrorHandler(
      entry.createErrorHandler(async (error, request, reply) => {
        reply.type('text/html').status((error as StatusError).status || 500)

        return renderer(
          request.url,
          Object.assign(
            {
              error,
            },
            reply.ctx,
          ),
        )
      }),
    )
  }
}

function createRenderer(
  _fastify: FastifyInstance,
  entry: Entry,
  render: RenderFunction | null | undefined,
  manifest: ViteManifest,
): Renderer {
  const files = manifest[`${entry.name}/client.ts`]
  const bundle = path.join('/', files.file)
  const preload = files.imports?.map((name) => path.join('/', manifest[name].file))
  const css = files.css?.map((name) => path.join('/', name))

  return (url, ctx) => {
    ctx = Object.assign({ url }, entry.ctx, ctx)

    const renderPromise = render?.(ctx)

    return entry
      .template({
        env: 'production',
        script: bundle,
        preload,
        css,
        content: renderPromise?.then((result) => result.html),
        head: renderPromise?.then((result) => result.head),
        state: renderPromise?.then((result) => result.state) || Promise.resolve(ctx),
      })
      .stream(true)
  }
}

function createHandler(renderer: Renderer): RouteHandler {
  return async function handler(request, reply) {
    reply.type('text/html')

    return renderer(request.url, reply.ctx)
  }
}

function createCachedHandler(cache: RenderCache): RouteHandler {
  return function cachedHandler(request, reply) {
    reply.type('text/html')

    return cache.get(request.url)
  }
}

async function buildCache(_entry: Entry, routes: Route[], renderer: Renderer, cache: RenderCache) {
  for (const route of routes) {
    if (route.cache) {
      let html = ''

      // TODO run preHandlers
      for await (const chunk of renderer(route.path) as Readable) {
        html += chunk
      }

      cache.set(route.path, html)
    }
  }
}