import { relative } from 'node:path'
import { normalizePath } from 'vite'
import type { Plugin } from 'vite'

/**
 * Tracks CSS files with `.inline.css` extension.
 * These CSS files will be inlined as <style> tags before scripts.
 *
 * Usage:
 *   import './layout.inline.css'
 *
 * The plugin records which source CSS files should be inlined. After the build,
 * getCriticalCSSOutputPaths() maps source paths to output paths
 * using the client manifest.
 */

const INLINE_CSS_EXT = '.inline.css'
// rust-side hook filter, so non-css ids never cross into js. exactly equivalent
// to the `id.endsWith(INLINE_CSS_EXT)` guard it replaced.
const INLINE_CSS_RE = /\.inline\.css$/

// source paths of CSS files with .inline.css extension (relative to root)
const criticalCSSSources = new Set<string>()

let root = ''

export function getCriticalCSSSources(): Set<string> {
  return criticalCSSSources
}

/**
 * Given the client manifest, returns the set of output CSS paths
 * that use .inline.css extension.
 */
export function getCriticalCSSOutputPaths(
  clientManifest: Record<string, { file: string; css?: string[] }>
): Set<string> {
  const outputPaths = new Set<string>()

  for (const [sourceKey, entry] of Object.entries(clientManifest)) {
    // direct CSS manifest entries: key is source path, entry.file is output
    if (sourceKey.endsWith('.css') && criticalCSSSources.has(sourceKey)) {
      outputPaths.add(`/${entry.file}`)
    }

    // also check if JS entries import critical CSS via their css array
    // by checking if the JS entry's source file imported any critical CSS
    // (this handles cssCodeSplit where CSS is grouped by chunk)
  }

  return outputPaths
}

export function criticalCSSPlugin(): Plugin {
  return {
    name: 'one:critical-css',
    enforce: 'pre',

    configResolved(config) {
      root = config.root
    },

    resolveId: {
      filter: { id: INLINE_CSS_RE },
      async handler(id, importer) {
        const resolved = await this.resolve(id, importer, { skipSelf: true })

        if (resolved) {
          // store as relative path to match manifest keys (forward slashes to match Vite)
          const relativePath = normalizePath(relative(root, resolved.id))
          criticalCSSSources.add(relativePath)
          return resolved
        }

        return null
      },
    },
  }
}
