import type { Assertion, ExpectPollOptions, PromisifyDomAssertion } from 'vitest'
import type { Locator } from 'vitest/browser'
import type { BrowserTraceEntryStatus } from './trace'
import { chai, expect } from 'vitest'
import { getType } from 'vitest/internal/browser'
import { getBrowserState, getWorkerState, now } from '../utils'
import { ariaMatchers } from './aria'
import { matchers } from './expect'
import { processTimeoutOptions } from './tester-utils'
import { recordBrowserTraceEntry } from './trace'

const kLocator = Symbol.for('$$vitest:locator')

function element<T extends HTMLElement | SVGElement | null | Locator>(elementOrLocator: T, options?: ExpectPollOptions): PromisifyDomAssertion<HTMLElement | SVGElement | null> {
  if (elementOrLocator != null && !(elementOrLocator instanceof HTMLElement) && !(elementOrLocator instanceof SVGElement) && !(kLocator in elementOrLocator)) {
    throw new Error(`Invalid element or locator: ${elementOrLocator}. Expected an instance of HTMLElement, SVGElement or Locator, received ${getType(elementOrLocator)}`)
  }

  const expectElement = expect.poll<HTMLElement | SVGElement | null>(function element(this: object) {
    if (elementOrLocator instanceof Element || elementOrLocator == null) {
      return elementOrLocator
    }

    const isNot = chai.util.flag(this, 'negate') as boolean
    const name = chai.util.flag(this, '_name') as string
    // special case for `toBeInTheDocument` matcher
    if (isNot && name === 'toBeInTheDocument') {
      return elementOrLocator.query()
    }
    if (name === 'toHaveLength') {
      // we know that `toHaveLength` requires multiple elements,
      // but types generally expect a single one
      return elementOrLocator.elements() as unknown as HTMLElement
    }

    if (name === 'toMatchScreenshot' && !chai.util.flag(this, '_poll.assert_once')) {
      // `toMatchScreenshot` should only run once after the element resolves
      chai.util.flag(this, '_poll.assert_once', true)
    }

    // element selector uses prettyDOM under the hood, which is an expensive call
    // that should not be called on each failed locator attempt to avoid memory leak:
    // https://github.com/vitest-dev/vitest/issues/7139
    const isLastPollAttempt = chai.util.flag(this, '_isLastPollAttempt')

    if (isLastPollAttempt) {
      return elementOrLocator.element()
    }

    const result = elementOrLocator.query()

    if (!result) {
      throw new Error(`Cannot find element with locator: ${JSON.stringify(elementOrLocator)}`)
    }

    return result
  }, processTimeoutOptions(options))

  chai.util.flag(expectElement, '_poll.element', true)

  // ask `expect.poll` to invoke trace after the assertion
  const currentTest = getWorkerState().current
  const hasActiveTrace = !!currentTest && getBrowserState().activeTraceTaskIds.has(currentTest.id)
  const hasActiveTraceView = !!currentTest && getBrowserState().browserTraceAttempts.has(currentTest.id)
  if (currentTest && (hasActiveTrace || hasActiveTraceView)) {
    const sourceError = new Error('__vitest_mark_trace__')
    const startTime = now()
    chai.util.flag(expectElement, '_poll.onSettled', async (meta: { assertion: Assertion; status: BrowserTraceEntryStatus }) => {
      const isNot = chai.util.flag(meta.assertion, 'negate')
      const name = chai.util.flag(meta.assertion, '_name') || '<unknown>'
      const baseName = `${isNot ? 'not.' : ''}${name}`
      const traceName = meta.status === 'fail' ? `${baseName} [ERROR]` : baseName
      const selector = !elementOrLocator || elementOrLocator instanceof Element
        ? undefined
        : elementOrLocator.serialize()
      if (hasActiveTraceView) {
        recordBrowserTraceEntry(currentTest, {
          name: traceName,
          kind: 'expect',
          status: meta.status,
          startTime,
          duration: now() - startTime,
          element: selector,
          stack: sourceError.stack,
        })
      }
      if (hasActiveTrace) {
        await getBrowserState().commands.triggerCommand(
          '__vitest_markTrace',
          [{
            name: traceName,
            element: selector,
            stack: sourceError.stack,
          }],
          sourceError,
        )
      }
    })
  }

  return expectElement
}

expect.extend(matchers)
expect.extend(ariaMatchers)
expect.element = element
