import {
  compile,
  env,
  Features,
  Instrumentation,
  normalizePath,
  optimize,
  toSourceMap,
} from '@tailwindcss/node'
import { clearRequireCache } from '@tailwindcss/node/require-cache'
import { Scanner } from '@tailwindcss/oxide'
import { realpathSync } from 'node:fs'
import fs from 'node:fs/promises'
import path from 'node:path'
import type {
  Environment,
  InternalResolveOptions,
  Plugin,
  ResolvedConfig,
  ViteDevServer,
} from 'vite'
import * as vite from 'vite'

const DEBUG = env.DEBUG
const SPECIAL_QUERY_RE = /[?&](?:worker|sharedworker|raw|url)\b/
const COMMON_JS_PROXY_RE = /\?commonjs-proxy/
const INLINE_STYLE_ID_RE = /[?&]index=\d+\.css$/

export type PluginOptions = {
  /**
   * Optimize and minify the output CSS.
   */
  optimize?: boolean | { minify?: boolean }
}

function createCustomResolver(
  resolvers: ((id: string, importer: string) => Promise<string | undefined>)[],
  filter = (_path: string) => true,
) {
  return async (id: string, base: string) => {
    // The resolver expects an `importer` file. We don't really know where the
    // current `id` was imported from, but Vite will essentially do a
    // `path.dirname(importer)` so it doesn't really matter.
    //
    // It does matter that this is a file, otherwise we would go up a directory,
    // which means that we would be resolving files from a parent folder first,
    // instead of the current folder we are in.
    let importer = path.resolve(base, '__placeholder__.css')

    for (let resolver of resolvers) {
      let resolved = await resolver(id, importer)

      // If we didn't resolve, we don't have to bail immediately, but we can try
      // the next resolver
      if (!resolved) continue

      if (resolved === id) continue

      // Looks like a relative file, let's resolve it to an absolute path
      if (resolved[0] === '.') resolved = path.resolve(base, resolved)

      // Must adhere to additional filters (e.g.: must be a .css file)
      if (!filter(resolved)) continue

      // If it's not an absolute path, then we don't really know how to read
      // the file from disk.
      if (!path.isAbsolute(resolved)) continue

      return resolved
    }
  }
}

export default function tailwindcss(opts: PluginOptions = {}): Plugin[] {
  let servers: ViteDevServer[] = []
  let config: ResolvedConfig | null = null
  let rootsByEnv = new DefaultMap<string, Map<string, Root>>((env: string) => new Map())

  let isSSR = false
  let shouldOptimize = true
  let minify = true

  function createRoot(env: Environment | null, id: string) {
    type ResolveFn = (id: string, base: string) => Promise<string | false | undefined>

    let customCssResolver: ResolveFn
    let customJsResolver: ResolveFn

    if (!env) {
      // Older, pre-environment Vite API
      // TODO: Can we drop this??
      let cssResolver = config!.createResolver({
        ...config!.resolve,
        extensions: ['.css'],
        mainFields: ['style'],
        conditions: ['style', 'development|production'],
        tryIndex: false,
        preferRelative: true,
      })

      let jsResolver = config!.createResolver(config!.resolve)

      customCssResolver = createCustomResolver(
        [
          (id, importer) => cssResolver(id, importer, true, isSSR),
          (id, importer) => cssResolver(id, importer, false, isSSR),
        ],
        (path) => path.endsWith('.css'),
      )
      customJsResolver = createCustomResolver(
        [
          (id, importer) => jsResolver(id, importer, true, isSSR),
          (id, importer) => jsResolver(id, importer, false, isSSR),
        ],
        (path) => !path.endsWith('.css'),
      )
    } else {
      type ResolveIdFn = (
        environment: Environment,
        id: string,
        importer?: string,
        aliasOnly?: boolean,
      ) => Promise<string | undefined>

      // There are cases where Environment API is available,
      // but `createResolver` is still overridden (for example astro v5)
      //
      // Copied as-is from vite, because this function is not a part of public API
      //
      // TODO: Remove this function and pre-environment code when Vite < 7 is no longer supported
      function createBackCompatIdResolver(
        config: ResolvedConfig,
        options?: Partial<InternalResolveOptions>,
      ): ResolveIdFn {
        const compatResolve = config.createResolver(options)
        let resolve: ResolveIdFn
        return async (environment, id, importer, aliasOnly) => {
          if (environment.name === 'client' || environment.name === 'ssr') {
            return compatResolve(id, importer, aliasOnly, environment.name === 'ssr')
          }
          resolve ??= vite.createIdResolver(config, options)
          return resolve(environment, id, importer, aliasOnly)
        }
      }

      // Newer Vite versions
      let cssResolver = createBackCompatIdResolver(env.config, {
        ...env.config.resolve,
        extensions: ['.css'],
        mainFields: ['style'],
        conditions: ['style', 'development|production'],
        tryIndex: false,
        preferRelative: true,
      })

      let jsResolver = createBackCompatIdResolver(env.config, env.config.resolve)

      customCssResolver = createCustomResolver(
        [
          (id, importer) => cssResolver(env, id, importer, true),
          (id, importer) => cssResolver(env, id, importer, false),
        ],
        (path) => path.endsWith('.css'),
      )
      customJsResolver = createCustomResolver(
        [
          (id, importer) => jsResolver(env, id, importer, true),
          (id, importer) => jsResolver(env, id, importer, false),
        ],
        (path) => !path.endsWith('.css'),
      )
    }

    return new Root(
      id,
      config!.root,
      // Currently, Vite only supports CSS source maps in development and they
      // are off by default. Check to see if we need them or not.
      config?.css.devSourcemap ?? false,
      customCssResolver,
      customJsResolver,
    )
  }

  return [
    {
      // Step 1: Scan source files for candidates
      name: '@tailwindcss/vite:scan',
      enforce: 'pre',

      configureServer(server) {
        servers.push(server)
      },

      async configResolved(_config) {
        config = _config
        isSSR = config.build.ssr !== false && config.build.ssr !== undefined

        // By default we optimize CSS during the build phase but if the user
        // provides explicit options we'll use those instead
        if (opts.optimize !== undefined) {
          shouldOptimize = opts.optimize !== false
        }

        // Minification is also performed when optimizing as long as it's also
        // enabled in Vite
        minify = shouldOptimize && config.build.cssMinify !== false

        // But again, the user can override that choice explicitly
        if (typeof opts.optimize === 'object') {
          minify = opts.optimize.minify !== false
        }
      },
    },

    {
      // Step 2 (serve mode): Generate CSS
      name: '@tailwindcss/vite:generate:serve',
      apply: 'serve',
      enforce: 'pre',
      transform: {
        filter: {
          id: {
            exclude: [/\/\.vite\//, SPECIAL_QUERY_RE, COMMON_JS_PROXY_RE],
            include: [/\.css(?:\?.*)?$/, /&lang\.css/, INLINE_STYLE_ID_RE],
          },
        },
        async handler(src, id) {
          if (!isPotentialCssRootFile(id)) return

          using I = new Instrumentation()
          DEBUG && I.start('[@tailwindcss/vite] Generate CSS (serve)')

          let roots = rootsByEnv.get(this.environment?.name ?? 'default')
          let root = roots.get(id)
          if (!root) {
            root ??= createRoot(this.environment ?? null, id)
            roots.set(id, root)
          }

          let result = await root.generate(src, (file) => this.addWatchFile(file), I)
          if (!result) {
            roots.delete(id)
            return src
          }

          DEBUG && I.end('[@tailwindcss/vite] Generate CSS (serve)')
          return result
        },
      },

      hotUpdate({ file, modules, timestamp, server }) {
        // Ensure full-reloads are triggered for files that are being watched by
        // Tailwind but aren't part of the module graph (like PHP or HTML
        // files). If we don't do this, then changes to those files won't
        // trigger a reload at all since Vite doesn't know about them.
        {
          // It's a little bit confusing, because due to the `addWatchFile`
          // calls, it _is_ part of the module graph but nothing is really
          // handling those files. These modules typically have an id of
          // undefined and/or have a type of 'asset'.
          //
          // If we call `addWatchFile` on a file that is part of the actual
          // module graph, then we will see a module for it with a type of `js`
          // and a type of `asset`. We are only interested if _all_ of them are
          // missing an id and/or have a type of 'asset', which is a strong
          // signal that the changed file is not being handled by Vite or any of
          // the plugins.
          //
          // Note: in Vite v7.0.6 the modules here will have a type of `js`, not
          // 'asset'. But it will also have a `HARD_INVALIDATED` state and will
          // do a full page reload already.
          //
          // Empty modules can be skipped since it means it's not `addWatchFile`d and thus irrelevant to Tailwind.
          let isExternalFile =
            modules.length > 0 &&
            modules.every((mod) => mod.type === 'asset' || mod.id === undefined)
          if (!isExternalFile) return

          // Skip if the module exists in other environments. SSR framework has
          // its own server side hmr/reload mechanism when handling server
          // only modules. See https://v6.vite.dev/guide/migration.html
          // > Updates to an SSR-only module no longer triggers a full page reload in the client. ...
          for (let environment of Object.values(server.environments)) {
            if (environment.name === this.environment.name) continue

            let modules = environment.moduleGraph.getModulesByFile(file)
            if (modules) {
              for (let module of modules) {
                if (module.type !== 'asset') {
                  return
                }
              }
            }
          }

          for (let env of new Set([this.environment.name, 'client'])) {
            let roots = rootsByEnv.get(env)
            if (roots.size === 0) continue

            // If the file is not being watched by any of the roots, then we can
            // skip the reload since it's not relevant to Tailwind CSS.
            if (!isScannedFile(file, modules, roots)) {
              continue
            }

            // https://vite.dev/changes/hotupdate-hook#migration-guide
            let invalidatedModules = new Set<vite.EnvironmentModuleNode>()
            for (let mod of modules) {
              this.environment.moduleGraph.invalidateModule(
                mod,
                invalidatedModules,
                timestamp,
                true,
              )
            }

            if (env === this.environment.name) {
              this.environment.hot.send({ type: 'full-reload' })
            } else if (server.hot.send) {
              server.hot.send({ type: 'full-reload' })
            } else if (server.ws.send) {
              server.ws.send({ type: 'full-reload' })
            }

            return []
          }
        }
      },
    },

    {
      // Step 2 (full build): Generate CSS
      name: '@tailwindcss/vite:generate:build',
      apply: 'build',
      enforce: 'pre',
      transform: {
        filter: {
          id: {
            exclude: [/\/\.vite\//, SPECIAL_QUERY_RE, COMMON_JS_PROXY_RE],
            include: [/\.css(?:\?.*)?$/, /&lang\.css/, INLINE_STYLE_ID_RE],
          },
        },
        async handler(src, id) {
          if (!isPotentialCssRootFile(id)) return

          using I = new Instrumentation()
          DEBUG && I.start('[@tailwindcss/vite] Generate CSS (build)')

          let roots = rootsByEnv.get(this.environment?.name ?? 'default')
          let root = roots.get(id)
          if (!root) {
            root ??= createRoot(this.environment ?? null, id)
            roots.set(id, root)
          }

          let result = await root.generate(src, (file) => this.addWatchFile(file), I)
          if (!result) {
            roots.delete(id)
            return src
          }
          DEBUG && I.end('[@tailwindcss/vite] Generate CSS (build)')

          if (shouldOptimize) {
            DEBUG && I.start('[@tailwindcss/vite] Optimize CSS')
            result = optimize(result.code, {
              minify,
              map: result.map,
            })
            DEBUG && I.end('[@tailwindcss/vite] Optimize CSS')
          }

          return result
        },
      },
    },
  ] satisfies Plugin[]
}

function getExtension(id: string) {
  let [filename] = id.split('?', 2)
  return path.extname(filename).slice(1)
}

function isPotentialCssRootFile(id: string) {
  if (id.includes('/.vite/')) return false

  // Don't intercept special static asset resources
  if (SPECIAL_QUERY_RE.test(id)) return false
  if (COMMON_JS_PROXY_RE.test(id)) return false

  let extension = getExtension(id)
  let isCssFile = extension === 'css' || id.includes('&lang.css') || id.match(INLINE_STYLE_ID_RE)

  return isCssFile
}

function idToPath(id: string) {
  return path.resolve(id.replace(/\?.*$/, ''))
}

/**
 * A Map that can generate default values for keys that don't exist.
 * Generated default values are added to the map to avoid recomputation.
 */
class DefaultMap<K, V> extends Map<K, V> {
  constructor(private factory: (key: K, self: DefaultMap<K, V>) => V) {
    super()
  }

  get(key: K): V {
    let value = super.get(key)

    if (value === undefined) {
      value = this.factory(key, this)
      this.set(key, value)
    }

    return value
  }
}

class Root {
  // The lazily-initialized Tailwind compiler components. These are persisted
  // throughout rebuilds but will be re-initialized if the rebuild strategy is
  // set to `full`.
  private compiler?: Awaited<ReturnType<typeof compile>>

  // The lazily-initialized Tailwind scanner.
  private scanner?: Scanner

  // List of all candidates that were being returned by the root scanner during
  // the lifetime of the root.
  private candidates: Set<string> = new Set<string>()

  // List of all build dependencies (e.g. imported  stylesheets or plugins) and
  // their last modification timestamp. If no mtime can be found, we need to
  // assume the file has always changed.
  private buildDependencies = new Map<string, number | null>()

  constructor(
    private id: string,
    private base: string,

    private enableSourceMaps: boolean,
    private customCssResolver: (id: string, base: string) => Promise<string | false | undefined>,
    private customJsResolver: (id: string, base: string) => Promise<string | false | undefined>,
  ) {}

  get scannedFiles() {
    return this.scanner?.files ?? []
  }

  // Generate the CSS for the root file. This can return false if the file is
  // not considered a Tailwind root. When this happened, the root can be GCed.
  public async generate(
    content: string,
    _addWatchFile: (file: string) => void,
    I: Instrumentation,
  ): Promise<
    | {
        code: string
        map: string | undefined
      }
    | false
  > {
    let inputPath = idToPath(this.id)

    function addWatchFile(file: string) {
      // Don't watch the input file since it's already a dependency and causes
      // issues with some setups (e.g. Qwik).
      if (file === inputPath) {
        return
      }

      // Scanning `.svg` file containing a `#` or `?` in the path will
      // crash Vite. We work around this for now by ignoring updates to them.
      //
      // https://github.com/tailwindlabs/tailwindcss/issues/16877
      if (/[#?].*\.svg$/.test(file)) {
        return
      }
      _addWatchFile(file)
    }

    let requiresBuildPromise = this.requiresBuild()
    let inputBase = path.dirname(path.resolve(inputPath))

    if (!this.compiler || !this.scanner || (await requiresBuildPromise)) {
      clearRequireCache(Array.from(this.buildDependencies.keys()))
      this.buildDependencies.clear()

      this.addBuildDependency(idToPath(inputPath))

      DEBUG && I.start('Setup compiler')
      let addBuildDependenciesPromises: Promise<void>[] = []
      this.compiler = await compile(content, {
        from: this.enableSourceMaps ? this.id : undefined,
        base: inputBase,
        shouldRewriteUrls: true,
        onDependency: (path) => {
          addWatchFile(path)
          addBuildDependenciesPromises.push(this.addBuildDependency(path))
        },

        customCssResolver: this.customCssResolver,
        customJsResolver: this.customJsResolver,
      })
      await Promise.all(addBuildDependenciesPromises)
      DEBUG && I.end('Setup compiler')

      DEBUG && I.start('Setup scanner')

      let sources = (() => {
        // Disable auto source detection
        if (this.compiler.root === 'none') {
          return []
        }

        // No root specified, auto-detect based on the `**/*` pattern
        if (this.compiler.root === null) {
          return [{ base: this.base, pattern: '**/*', negated: false }]
        }

        // Use the specified root
        return [{ ...this.compiler.root, negated: false }]
      })().concat(this.compiler.sources)

      this.scanner = new Scanner({ sources })
      DEBUG && I.end('Setup scanner')
    } else {
      for (let buildDependency of this.buildDependencies.keys()) {
        addWatchFile(buildDependency)
      }
    }

    if (
      !(
        this.compiler.features &
        (Features.AtApply |
          Features.JsPluginCompat |
          Features.ThemeFunction |
          Features.Utilities |
          Features.Variants)
      )
    ) {
      return false
    }

    if (this.compiler.features & Features.Utilities) {
      // This should not be here, but right now the Vite plugin is setup where we
      // setup a new scanner and compiler every time we request the CSS file
      // (regardless whether it actually changed or not).
      DEBUG && I.start('Scan for candidates')
      for (let candidate of this.scanner.scan()) {
        this.candidates.add(candidate)
      }
      DEBUG && I.end('Scan for candidates')
    }

    if (this.compiler.features & Features.Utilities) {
      DEBUG && I.start('Register dependency messages')
      // Watch individual files found via custom `@source` paths
      for (let file of this.scanner.files) {
        addWatchFile(file)
      }

      // Watch globs found via custom `@source` paths
      for (let glob of this.scanner.globs) {
        if (glob.pattern[0] === '!') continue

        let relative = path.relative(this.base, glob.base)
        if (relative[0] !== '.') {
          relative = './' + relative
        }
        // Ensure relative is a posix style path since we will merge it with the
        // glob.
        relative = normalizePath(relative)

        addWatchFile(path.posix.join(relative, glob.pattern))

        let root = this.compiler.root

        if (root !== 'none' && root !== null) {
          let basePath = normalizePath(path.resolve(root.base, root.pattern))

          let isDir = await fs.stat(basePath).then(
            (stats) => stats.isDirectory(),
            () => false,
          )

          if (!isDir) {
            throw new Error(
              `The path given to \`source(…)\` must be a directory but got \`source(${basePath})\` instead.`,
            )
          }
        }
      }
      DEBUG && I.end('Register dependency messages')
    }

    DEBUG && I.start('Build CSS')
    let code = this.compiler.build([...this.candidates])
    DEBUG && I.end('Build CSS')

    DEBUG && I.start('Build Source Map')
    let map = this.enableSourceMaps ? toSourceMap(this.compiler.buildSourceMap()).raw : undefined
    DEBUG && I.end('Build Source Map')

    return {
      code,
      map,
    }
  }

  private async addBuildDependency(path: string) {
    let mtime: number | null = null
    try {
      mtime = (await fs.stat(path)).mtimeMs
    } catch {}
    this.buildDependencies.set(path, mtime)
  }

  private async requiresBuild(): Promise<boolean> {
    for (let [path, mtime] of this.buildDependencies) {
      if (mtime === null) return true
      try {
        let stat = await fs.stat(path)
        if (stat.mtimeMs > mtime) {
          return true
        }
      } catch {
        return true
      }
    }
    return false
  }
}

function isScannedFile(
  file: string,
  modules: vite.EnvironmentModuleNode[],
  roots: Map<string, Root>,
) {
  let seen = new Set()
  let q = [...modules]
  let checks = {
    file,
    get realpath() {
      try {
        let realpath = realpathSync(file)
        Object.defineProperty(checks, 'realpath', { value: realpath })
        return realpath
      } catch {
        return null
      }
    },
  }

  while (q.length > 0) {
    let module = q.shift()!
    if (seen.has(module)) continue
    seen.add(module)

    if (module.id) {
      let root = roots.get(module.id)

      if (root) {
        // If the file is part of the scanned files for this root, then we know
        // for sure that it's being watched by any of the Tailwind CSS roots. It
        // doesn't matter which root it is since it's only used to know whether
        // we should trigger a full reload or not.
        if (
          root.scannedFiles.includes(checks.file) ||
          (checks.realpath && root.scannedFiles.includes(checks.realpath))
        ) {
          return true
        }
      }
    }

    // Keep walking up the tree until we find a root.
    for (let importer of module.importers) {
      q.push(importer)
    }
  }

  return false
}