build-extension.ts

import type { OutputOptions } from 'rollup'
import type * as vite from 'vite'
import type * as wxt from 'wxt'
import { writeFile } from 'node:fs/promises'
import path from 'node:path'
import { parseHTML } from 'linkedom'
import { hash } from 'ohash'
import {
  addViteConfig,
  defineWxtModule,
} from 'wxt/modules'

type AddedViteConfig = ReturnType<Parameters<typeof addViteConfig>[1]>
type AddedVitePlugins = NonNullable<NonNullable<AddedViteConfig>['plugins']>

export default defineWxtModule({
  async setup(wxt) {
    wxt.config.alias[`/src/main.ts`] = `./src/main.ts`
    wxt.config.alias[`/src/sidepanel.ts`] = `./src/sidepanel.ts`
    wxt.config.manifest.options_page = `options.html`
    wxt.hook(`entrypoints:grouped`, (_, groups) => {
      groups.push([{
        type: `options`,
        name: `options`,
        options: { openInTab: true },
        inputPath: path.resolve(wxt.config.root, `./index.html`),
        outputDir: wxt.config.outDir,
        skipped: false,
      }])
      groups.push([{
        type: `sidepanel`,
        name: `sidepanel`,
        options: { openAtInstall: true, browserStyle: true },
        inputPath: path.resolve(wxt.config.root, `./index.html`),
        outputDir: wxt.config.outDir,
        skipped: false,
      }])
    })
    wxt.hook(`vite:build:extendConfig`, (_, config) => {
      if (config.build?.rollupOptions?.input && config.build?.rollupOptions?.output) {
        const input = config.build?.rollupOptions.input as Record<string, string>
        const wxtOutput = config.build?.rollupOptions.output as OutputOptions
        if (input.options || input.sidepanel) {
          wxtOutput.manualChunks = (id) => {
            if (id.includes(`node_modules`)) {
              if (id.includes(`prettier`))
                return `prettier`
              if (id.includes(`katex`))
                return `katex`
              if (id.includes(`mermaid`) || id.includes(`highlight.js`))
                return `mermaid-vendors`
            }
          }
        }
      }
    })
    addViteConfig(wxt, () => ({
      plugins: toWxtPluginOptions([
        htmlScriptToVirtual(wxt.config, () => wxt.server),
        vueDevtoolsHack(wxt.config, () => wxt.server),
        wxt.config.command === `build`
          ? htmlScriptToLocal(wxt)
          : undefined,
      ]),
    }))
  },
})

// Stored outside the plugin to effect all instances of the htmlScriptToVirtual plugin.
const inlineScriptContents: Record<string, string> = {}
const SCRIPT_FILE_NAME_REGEX = /\/([^/]+)\.js$/

function isDefined<T>(value: T | undefined): value is T {
  return value !== undefined
}

function toWxtPluginOptions(
  plugins: ReadonlyArray<vite.PluginOption | undefined>,
): AddedVitePlugins {
  return plugins.filter(isDefined) as unknown as AddedVitePlugins
}

export function htmlScriptToVirtual(
  config: wxt.ResolvedConfig,
  getWxtDevServer: () => wxt.WxtDevServer | undefined,
): vite.PluginOption {
  const virtualInlineScript = `virtual:md-inline-script`
  const resolvedVirtualInlineScript = `\0${virtualInlineScript}`

  const server = getWxtDevServer?.()
  return [
    {
      name: `md:dev-html-prerender`,
      apply: `build`,
      transformIndexHtml: {
        order: `post`,
        async handler(html) {
          if (server == null) {
            return html
          }
          const { document } = parseHTML(html)
          // Replace inline script with virtual module served via dev server.
          // Extension CSP blocks inline scripts, so that's why we're pulling them out.
          const promises: Promise<void>[] = []
          const inlineScripts = document.querySelectorAll(`script[src^=http]`)
          inlineScripts.forEach((script) => {
            const url = script.getAttribute(`src`) ?? ``
            if (url?.startsWith(`http://localhost`)) {
              return
            }
            const p = doFetch(url).then((textContent) => {
              const key = hash(textContent)
              inlineScriptContents[key] = textContent
              script.setAttribute(`src`, `${server.origin}/@id/${virtualInlineScript}?${key}`)
              if (script.hasAttribute(`id`)) {
                script.setAttribute(`type`, `module`)
              }
            })
            promises.push(p)
          })
          await Promise.all(promises)
          const newHtml = document.toString()
          config.logger.debug(`\nhtmlScriptToVirtual Old HTML:\n${html}`)
          config.logger.debug(`\nhtmlScriptToVirtual New HTML:\n${newHtml}`)
          return newHtml
        },
      },
    },
    {
      name: `md:virtualize-react-refresh`,
      apply: `serve`,
      resolveId(id) {
        // Resolve inline scripts
        if (id.startsWith(virtualInlineScript)) {
          return `\0${id}`
        }

        // Ignore chunks during HTML file pre-rendering
        if (id.startsWith(`/chunks/`)) {
          return `\0noop`
        }
      },
      load(id) {
        // Resolve virtualized inline scripts
        if (id.startsWith(resolvedVirtualInlineScript)) {
          // id="virtual:md-inline-script?<hash>"
          const key = id.substring(id.indexOf(`?`) + 1)
          return inlineScriptContents[key]
        }

        // Ignore chunks during HTML file pre-rendering
        if (id === `\0noop`) {
          return ``
        }
      },
    },
  ]
}

export function htmlScriptToLocal(
  wxt: wxt.Wxt,
): vite.Plugin {
  return {
    name: `md:build-html-prerender`,
    apply: `build`,
    transformIndexHtml: {
      order: `post`,
      async handler(html) {
        const { document } = parseHTML(html)
        const promises: Promise<void>[] = []
        const httpScripts = document.querySelectorAll(`script[src^=http]`)
        if (httpScripts.length > 0) {
          httpScripts.forEach((script) => {
            const url = script.getAttribute(`src`) ?? ``
            if (url?.startsWith(`http://localhost`)) {
              return
            }
            const p = (async () => {
              const textContent = await doFetch(url)
              const key = hash(textContent)
              let jsName = url.match(SCRIPT_FILE_NAME_REGEX)?.[1] ?? `.js`
              if (url.indexOf(`?`) > 0) {
                jsName = `${url.substring(url.indexOf(`?`) + 1)}.js`
              }
              const fileName = `${jsName.split(`.`)[0]}-${key}.js`
              // write to file
              const outFile = path.resolve(wxt.config.outDir, `./${fileName}`)
              await writeFile(outFile, textContent, `utf8`)
              script.setAttribute(`src`, `/${fileName}`)
            })()
            promises.push(p)
          })
        }

        // Replace inline script with virtual module served via dev server.
        // Extension CSP blocks inline scripts, so that's why we're pulling them
        // out.
        const inlineScripts = document.querySelectorAll(`script:not([src])`)
        if (inlineScripts.length > 0) {
          inlineScripts.forEach((script) => {
            const p = (async () => {
              // Save the textContent for later
              const textContent = script.textContent ?? ``
              const key = hash(textContent)
              const fileName = `md-inline-${key}.js`
              // write to file
              const outFile = path.resolve(wxt.config.outDir, `./${fileName}`)
              await writeFile(outFile, textContent, `utf8`)
              // Replace unsafe inline script
              const virtualScript = document.createElement(`script`)
              virtualScript.src = `/${fileName}`
              script.replaceWith(virtualScript)
            })()
            promises.push(p)
          })
        }
        await Promise.all(promises)
        const newHtml = document.toString()
        wxt.config.logger.debug(`\nhtmlScriptToLocal Old HTML:\n${html}`)
        wxt.config.logger.debug(`\nhtmlScriptToLocal New HTML:\n${newHtml}`)
        return newHtml
      },
    },
  }
}
export function vueDevtoolsHack(
  config: wxt.ResolvedConfig,
  getWxtDevServer: () => wxt.WxtDevServer | undefined,
): vite.Plugin {
  const server = getWxtDevServer?.()
  return {
    name: `md:vue-devtools-hack`,
    apply: `build`,
    transformIndexHtml: {
      order: `post`,
      handler(html) {
        const { document } = parseHTML(html)
        const inlineScripts = document.querySelectorAll(`script[src^='/@id/virtual:']`)
        inlineScripts.forEach((script) => {
          const src = script.getAttribute(`src`)
          const newSrc = `${server?.origin}${src}`
          script.setAttribute(`src`, newSrc)
        })
        const newHtml = document.toString()
        config.logger.debug(`Old HTML:\n${html}`)
        config.logger.debug(`New HTML:\n${newHtml}`)

        return newHtml
      },
    },
  }
}

async function doFetch(
  url: string,
): Promise<string> {
  let content: string = ``
  const res = await fetch(url)
  if (res.status < 300) {
    content = await res.text()
  }
  else {
    throw new Error(
      `Failed to fetch "${url}". `,
    )
  }
  return content
}