import MagicString from 'magic-string'
import { createFilter } from '@rollup/pluginutils'
import type { Plugin, OutputAsset, RenderedChunk } from 'rollup'
import type { Directive, ExpressionStatement } from 'estree'

type EnsureDirectivesOptions = {
  include?: string[]
  exclude?: string[]
}

/**
 * This is a plugin that ensures directives like "use client" gets put back at the top of files after bundling.
 * In order for this plugin to work, every input file needs to be outputed in a separate output chunk. Otherwise,
 * the directive would not be scoped to the related file.
 * This can either be done by using the output.preserveModules or by using glob() as stated in Rollup documentation:
 * https://rollupjs.org/configuration-options/#input
 *
 * @param {Object} options - Plugin options
 * @param {string[]} options.include - Minimatch patterns to include in parsing
 * @param {string[]} options.exclude - Minimatch patterns to skip from parsing
 *
 */

function isDirective(body: Directive | ExpressionStatement): body is Directive {
  return (body as Directive).directive !== undefined
}

function isChunk(chunk: OutputAsset | RenderedChunk): chunk is RenderedChunk {
  return (chunk as RenderedChunk).modules !== undefined
}

export default function ensureDirectives({
  include = [],
  exclude = [],
}: EnsureDirectivesOptions = {}): Plugin {
  // Skip CSS files by default, as this.parse() does not work on them
  const excludePatterns = ['**/*.css', ...exclude]
  const includePatterns = include
  const filter = createFilter(includePatterns, excludePatterns)

  return {
    name: 'ensure-directives',
    // Capture directives metadata during the transform phase
    transform(code, id) {
      // Skip files that are excluded or that are implicitly excluded by the include pattern
      if (!filter(id)) return

      const ast = this.parse(code)

      if (ast.type === 'Program' && ast.body) {
        const filteredBodies = ast.body.filter(Boolean)

        const directives = filteredBodies.reduce<string[]>(
          (acc, filteredBody) => {
            if (
              filteredBody &&
              filteredBody.type === 'ExpressionStatement' &&
              isDirective(filteredBody)
            ) {
              acc.push(filteredBody.directive)
            }
            return acc
          },
          [],
        )
        // Map is set to null since the sourceMap is left untouched
        if (directives.length) {
          return {
            code,
            ast,
            map: null,
            meta: { ensureDirectives: directives },
          }
        }
      }
      return { code, ast, map: null }
    },
    renderChunk: {
      order: 'post',
      handler(code, chunk) {
        // Only do this for RenderedChunk, not OutputAssets.
        if (isChunk(chunk)) {
          const modulesKeys = Object.keys(chunk.modules)

          for (const moduleId of modulesKeys) {
            const directives: string[] | false =
              this.getModuleInfo(moduleId)?.meta?.ensureDirectives

            if (directives) {
              const directiveStrings = directives
                .map((directive) => `"${directive}"`)
                .join(';\n')

              const s = new MagicString(code)
              s.prepend(`${directiveStrings};\n`)
              const srcMap = s.generateMap({ includeContent: true })

              return { code: s.toString(), map: srcMap }
            }
          }
        }
        return null
      },
    },
  }
}