import { promises as fs } from "fs"
import path from "path"
import zlib from "zlib"
import { promisify } from "util"
import type { Plugin, ResolvedConfig } from "vite"

const gzip = promisify(zlib.gzip)
const brotliCompress = promisify(zlib.brotliCompress)
const compressibleFileRE = /\.(js|mjs|json|css|html)$/i
const defaultThreshold = 1025

type CompressionKind = "gzip" | "brotli"

const compressionHandlers: Record<
	CompressionKind,
	{ ext: string; compress: (content: Buffer) => Promise<Buffer> }
> = {
	gzip: {
		ext: ".gz",
		compress: (content) => gzip(content, { level: zlib.constants.Z_BEST_COMPRESSION })
	},
	brotli: {
		ext: ".br",
		compress: (content) =>
			brotliCompress(content, {
				params: {
					[zlib.constants.BROTLI_PARAM_QUALITY]: zlib.constants.BROTLI_MAX_QUALITY,
					[zlib.constants.BROTLI_PARAM_MODE]: zlib.constants.BROTLI_MODE_TEXT
				}
			})
	}
}

async function collectFiles(rootDir: string): Promise<string[]> {
	const entries = await fs.readdir(rootDir, { withFileTypes: true })
	const files = await Promise.all(
		entries.map(async (entry) => {
			const fullPath = path.join(rootDir, entry.name)
			if (entry.isDirectory()) {
				return collectFiles(fullPath)
			}
			return compressibleFileRE.test(entry.name) ? [fullPath] : []
		})
	)
	return files.flat()
}

function createCompressionPlugin(kind: CompressionKind): Plugin {
	const handler = compressionHandlers[kind]
	let config: ResolvedConfig | undefined

	return {
		name: `local:compression:${kind}`,
		apply: "build",
		enforce: "post",
		configResolved(resolvedConfig) {
			config = resolvedConfig
		},
		async closeBundle() {
			const outputDir = path.resolve(process.cwd(), config?.build.outDir ?? "dist")
			const files = await collectFiles(outputDir)
			const compressedEntries: Array<{ file: string; originalKb: string; compressedKb: string }> =
				[]

			await Promise.all(
				files.map(async (filePath) => {
					const stat = await fs.stat(filePath)
					if (stat.size < defaultThreshold) {
						return
					}

					const content = await fs.readFile(filePath)
					const compressed = await handler.compress(content)
					const outputFile = `${filePath}${handler.ext}`

					await fs.writeFile(outputFile, compressed)
					compressedEntries.push({
						file: path.relative(outputDir, outputFile).replace(/\\/g, "/"),
						originalKb: (stat.size / 1024).toFixed(2),
						compressedKb: (compressed.byteLength / 1024).toFixed(2)
					})
				})
			)

			if (!compressedEntries.length) {
				return
			}

			compressedEntries.sort((a, b) => a.file.localeCompare(b.file))
			config?.logger.info(`\n[compression:${kind}] generated ${compressedEntries.length} files`)
			for (const entry of compressedEntries) {
				config?.logger.info(
					`${path.basename(outputDir)}/${entry.file} ${entry.originalKb}kb -> ${entry.compressedKb}kb`
				)
			}
			config?.logger.info("")
		}
	}
}

export default (env: Record<string, string>) => {
	const { VITE_BUILD_COMPRESS } = env
	const plugins: Plugin[] = []
	if (!VITE_BUILD_COMPRESS) {
		return plugins
	}

	const compressionList = VITE_BUILD_COMPRESS.split(",").map((item) =>
		item.trim()
	) as CompressionKind[]
	if (compressionList.includes("gzip")) {
		plugins.push(createCompressionPlugin("gzip"))
	}
	if (compressionList.includes("brotli")) {
		plugins.push(createCompressionPlugin("brotli"))
	}
	return plugins
}