import type { SerializedLocator } from '@vitest/browser'
import type { ScreenshotOptions } from 'vitest/browser'
import type { BrowserCommandContext } from 'vitest/node'
import { mkdir } from 'node:fs/promises'
import {
  assertBrowserApiWrite,
  assertBrowserFileAccess,
  resolveScreenshotPath,
} from '@vitest/browser'
import { dirname, normalize } from 'pathe'
import { getDescribedLocator } from './utils'

interface ScreenshotCommandOptions extends Omit<ScreenshotOptions, 'element' | 'mask'> {
  element?: SerializedLocator
  mask?: readonly SerializedLocator[]
  target?: 'element' | 'page'
}

const SCREENSHOT_STYLES = /* css */ `
  iframe[data-vitest="true"] {
    position: absolute !important;
    inset: 0 !important;
    z-index: ${Number.MAX_SAFE_INTEGER} !important;
    transform: none !important;
  }
`

/**
 * Takes a screenshot using the provided browser context and returns a buffer and the expected screenshot path.
 *
 * **Note**: the returned `path` indicates where the screenshot *might* be found.
 * It is not guaranteed to exist, especially if `options.save` is `false`.
 *
 * @throws {Error} If the function is not called within a test or if the browser provider does not support screenshots.
 */
export async function takeScreenshot(
  context: BrowserCommandContext,
  name: string,
  options: Omit<ScreenshotCommandOptions, 'base64'>,
): Promise<{ buffer: Buffer<ArrayBufferLike>; path: string }> {
  if (!context.testPath) {
    throw new Error(`Cannot take a screenshot without a test path`)
  }

  const path = resolveScreenshotPath(context.testPath, name, context.project.config, options.path)

  // playwright does not need a screenshot path if we don't intend to save it
  let savePath: string | undefined

  if (options.save) {
    savePath = normalize(path)

    assertBrowserApiWrite(context.project, savePath)
    assertBrowserFileAccess(context.project, savePath)

    await mkdir(dirname(savePath), { recursive: true })
  }

  const mask = options.mask?.map((selector) => getDescribedLocator(context, selector))
  const style = context.project.config.browser.ui
    ? options.style === undefined
      ? SCREENSHOT_STYLES
      : SCREENSHOT_STYLES + options.style
    : options.style

  if (options.element) {
    const { element: selector, target: _target, ...config } = options
    const element = getDescribedLocator(context, selector)
    const buffer = await element.screenshot({
      ...config,
      mask,
      path: savePath,
      style,
    })
    return { buffer, path }
  }

  const { target, ...config } = options
  const buffer =
    target === 'page'
      ? await context.page.screenshot({
          ...config,
          mask,
          path: savePath,
          style,
        })
      : await getDescribedLocator(context, {
          selector: 'body',
          locator: "locator('body')",
        }).screenshot({
          ...config,
          mask,
          path: savePath,
          style,
        })
  return { buffer, path }
}