import type { PreviewServer, ViteDevServer } from "vite";

/**
 * Vite server hook for configuring cache control headers
 * @param server - Vite dev server or preview server instance
 */
export function cacheServerHook<T extends ViteDevServer | PreviewServer>(
  server: T,
) {
  server.middlewares.use(async (request, response, next) => {
    let cacheControlValue = "public, max-age=86400, must-revalidate";
    if (request.url?.startsWith("/assets/")) {
      cacheControlValue = "public, max-age=31536000, immutable";
    } else if (
      request.url === "/" ||
      request.url?.startsWith("/?") ||
      request.url?.endsWith(".html")
    ) {
      cacheControlValue = "no-cache";
    }

    response.setHeader("Cache-Control", cacheControlValue);

    next();
  });
}