import type { FileSpecification, VitestRunner } from './types/runner'
import type { File, SuiteHooks } from './types/tasks'
import { processError } from '@vitest/utils/error' // TODO: load dynamically
import { toArray } from '@vitest/utils/helpers'
import { collectorContext, setFileContext } from './context'
import { getHooks, setHooks } from './map'
import { runSetupFiles } from './setup'
import {
  clearCollectorContext,
  createSuiteHooks,
  getDefaultSuite,
} from './suite'
import {
  calculateSuiteHash,
  createFileTask,
  interpretTaskModes,
  someTasksAreOnly,
} from './utils/collect'

const now = globalThis.performance ? globalThis.performance.now.bind(globalThis.performance) : Date.now

export async function collectTests(
  specs: string[] | FileSpecification[],
  runner: VitestRunner,
): Promise<File[]> {
  const files: File[] = []

  const config = runner.config
  const $ = runner.trace!

  for (const spec of specs) {
    const filepath = typeof spec === 'string' ? spec : spec.filepath
    await $(
      'collect_spec',
      { 'code.file.path': filepath },
      async () => {
        const testLocations = typeof spec === 'string' ? undefined : spec.testLocations

        const file = createFileTask(filepath, config.root, config.name, runner.pool, runner.viteEnvironment)
        setFileContext(file, Object.create(null))
        file.shuffle = config.sequence.shuffle

        runner.onCollectStart?.(file)

        clearCollectorContext(file, runner)

        try {
          const setupFiles = toArray(config.setupFiles)
          if (setupFiles.length) {
            const setupStart = now()
            await runSetupFiles(config, setupFiles, runner)
            const setupEnd = now()
            file.setupDuration = setupEnd - setupStart
          }
          else {
            file.setupDuration = 0
          }

          const collectStart = now()

          await runner.importFile(filepath, 'collect')

          const durations = runner.getImportDurations?.()
          if (durations) {
            file.importDurations = durations
          }

          const defaultTasks = await getDefaultSuite().collect(file)

          const fileHooks = createSuiteHooks()
          mergeHooks(fileHooks, getHooks(defaultTasks))

          for (const c of [...defaultTasks.tasks, ...collectorContext.tasks]) {
            if (c.type === 'test' || c.type === 'suite') {
              file.tasks.push(c)
            }
            else if (c.type === 'collector') {
              const suite = await c.collect(file)
              if (suite.name || suite.tasks.length) {
                mergeHooks(fileHooks, getHooks(suite))
                file.tasks.push(suite)
              }
            }
            else {
              // check that types are exhausted
              c satisfies never
            }
          }

          setHooks(file, fileHooks)
          file.collectDuration = now() - collectStart
        }
        catch (e) {
          const error = processError(e)
          file.result = {
            state: 'fail',
            errors: [error],
          }

          const durations = runner.getImportDurations?.()
          if (durations) {
            file.importDurations = durations
          }
        }

        calculateSuiteHash(file)

        const hasOnlyTasks = someTasksAreOnly(file)
        interpretTaskModes(
          file,
          config.testNamePattern,
          testLocations,
          hasOnlyTasks,
          false,
          config.allowOnly,
        )

        if (file.mode === 'queued') {
          file.mode = 'run'
        }

        files.push(file)
      },
    )
  }

  return files
}

function mergeHooks(baseHooks: SuiteHooks, hooks: SuiteHooks): SuiteHooks {
  for (const _key in hooks) {
    const key = _key as keyof SuiteHooks
    baseHooks[key].push(...(hooks[key] as any))
  }

  return baseHooks
}
