import {
  bold,
  red
} from 'colorette'
import {
  AsyncZippable,
  zip
} from 'fflate'
import { filesize } from 'filesize'
import {
  readdir,
  readFile,
  writeFile
} from 'node:fs/promises'
import { promisify } from 'node:util'
import { Plugin } from 'vite'
import { relative } from 'node:path'


async function* traversalDir(path: string): AsyncGenerator<[ string, string, Uint8Array ]> {
  yield [ 'startDir', path, new Uint8Array(1) ]
  let dirents = await readdir(path, {withFileTypes: true})
  for (let dirent of dirents) {
    let subDir = `${path}/${dirent.name}`
    if (dirent.isDirectory()) {
      yield* traversalDir(subDir)
    } else {
      yield  [ 'file', subDir, await readFile(subDir) ]
    }
  }
  yield   [ 'endDir', path, new Uint8Array(1) ]
}

async function doZip(path: string) {
  let data: AsyncZippable = {}
  for await (const e of traversalDir(path)) {
    if (e[0] === 'file') {
      data[relative(path, e[1])] = e[2]
    }
  }
  return await promisify(zip)(data)
}

export default function zipDist(name?: string, addVersion: boolean = false): Plugin {
  if (!!name && !name.endsWith('.zip')) {
    name = `${name}.zip`
  } else if (!name) {
    let {npm_package_name = 'bundle', npm_package_version} = process.env
    name = !!npm_package_version && addVersion ?
        `${npm_package_name}-${npm_package_version}.zip` : `${npm_package_name}.zip`
  }

  return {
    name: 'vite-plugin-dist-zip',
    writeBundle: {
      sequential: true,
      order: 'post',
      handler(options, bundle) {
        doZip(options.dir!)
            .then(value => Promise.all([ writeFile(`./dist/${name}`, value), Promise.resolve(value.byteLength) ]))
            .then(value => console.log(red(`\n压缩完成 ${bold(`dist/${name}`)}    ${filesize(value[1])}`)))
      }
    }
  }
}

