import type { CoverageMap } from '@vitest/istanbul-lib-coverage'
import type { Instrumenter } from '@vitest/istanbul-lib-instrument'
import type { ProxifiedModule } from 'magicast'
import type { CoverageProvider, ReportContext, Vite, Vitest } from 'vitest/node'
import { existsSync, promises as fs } from 'node:fs'
import { addMapping, GenMapping, toEncodedMap } from '@jridgewell/gen-mapping'
import { eachMapping, TraceMap } from '@jridgewell/trace-mapping'
import * as libCoverage from '@vitest/istanbul-lib-coverage'
import { createInstrumenter } from '@vitest/istanbul-lib-instrument'
import * as libReport from '@vitest/istanbul-lib-report'
import * as libSourceMaps from '@vitest/istanbul-lib-source-maps'
import { parseModule } from 'magicast'
import { createDebug } from 'obug'
import c from 'tinyrainbow'
import { BaseCoverageProvider, isCSSRequest } from 'vitest/node'
import { version } from '../package.json' with { type: 'json' }
import { commands } from './commands'
import { COVERAGE_STORE_KEY } from './constants'

const debug = createDebug('vitest:coverage')

export class IstanbulCoverageProvider extends BaseCoverageProvider implements CoverageProvider {
  name = 'istanbul' as const
  version: string = version
  instrumenter!: Instrumenter

  private transformedModuleIds = new Set<string>()

  initialize(ctx: Vitest): void {
    this._initialize(ctx)

    for (const project of ctx.projects) {
      if (project.isBrowserEnabled() && project.browser) {
        for (const [name, command] of Object.entries(commands)) {
          project.browser.registerCommand(`__vitest_${name}` as any, command)
        }
      }
    }

    if (this.options.instrumenter) {
      this.instrumenter = this.options.instrumenter({
        coverageVariable: COVERAGE_STORE_KEY,
        coverageGlobalScope: 'globalThis',
        coverageGlobalScopeFunc: false,
        ignoreClassMethods: this.options.ignoreClassMethods,
      }) as Instrumenter
    } else {
      this.instrumenter = createInstrumenter({
        produceSourceMap: true,
        autoWrap: false,
        esModules: true,
        compact: false,
        coverageVariable: COVERAGE_STORE_KEY,
        coverageGlobalScope: 'globalThis',
        coverageGlobalScopeFunc: false,
        ignoreClassMethods: this.options.ignoreClassMethods,
        ignoreLines: true,
      })
    }
  }

  requiresTransform(id: string): boolean {
    // Istanbul/babel cannot instrument CSS - e.g. Vue imports end up here.
    // File extension itself is .vue, but it contains CSS.
    // e.g. "Example.vue?vue&type=style&index=0&scoped=f7f04e08&lang.css"
    if (isCSSRequest(id)) {
      return false
    }

    if (!this.isIncluded(removeQueryParameters(id))) {
      return false
    }

    return true
  }

  onFileTransform(
    sourceCode: string,
    id: string,
    pluginCtx: Vite.Rollup.TransformPluginContext,
  ): { code: string; map: any } | undefined {
    if (!this.requiresTransform(id)) {
      return
    }

    const sourceMap = pluginCtx.getCombinedSourcemap()
    sourceMap.sources = sourceMap.sources.map(removeQueryParameters)

    sourceCode = sourceCode
      // Exclude SWC's decorators that are left in source maps
      .replaceAll('_ts_decorate', '/* istanbul ignore next */_ts_decorate')

      // Exclude in-source test's test cases
      .replaceAll(/(if +\(import\.meta\.vitest\))/g, '/* istanbul ignore next */ $1')

    const code = this.instrumenter.instrumentSync(sourceCode, id, sourceMap as any)

    if (!id.includes('vitest-uncovered-coverage=true')) {
      const transformMap = new GenMapping(sourceMap)

      eachMapping(new TraceMap(sourceMap as any), (mapping) => {
        addMapping(transformMap, {
          generated: { line: mapping.generatedLine, column: mapping.generatedColumn },
          original: { line: mapping.generatedLine, column: mapping.generatedColumn },
          content: sourceCode,
          name: mapping.name || '',
          source: mapping.source || '',
        })
      })

      const encodedMap = toEncodedMap(transformMap)
      delete encodedMap.file
      delete encodedMap.ignoreList
      delete encodedMap.sourceRoot

      this.instrumenter.instrumentSync(sourceCode, id, encodedMap as any)
    }

    const map = this.instrumenter.lastSourceMap() as any
    this.transformedModuleIds.add(id)

    return { code, map }
  }

  createCoverageMap(): libCoverage.CoverageMap {
    return libCoverage.createCoverageMap({})
  }

  async generateCoverage({ allTestsRun }: ReportContext): Promise<CoverageMap> {
    const start = debug.enabled ? performance.now() : 0

    const coverageMap = this.createCoverageMap()
    let coverageMapByEnvironment = this.createCoverageMap()

    await this.readCoverageFiles<CoverageMap>({
      onFileRead(coverage) {
        coverageMapByEnvironment.merge(coverage)
      },
      onFinished: async () => {
        // Source maps can change based on projectName and transform mode.
        // Coverage transform re-uses source maps so we need to separate transforms from each other.
        const transformedCoverage = await transformCoverage(coverageMapByEnvironment)
        coverageMap.merge(transformedCoverage)

        coverageMapByEnvironment = this.createCoverageMap()
      },
      onDebug: debug,
    })

    // Include untested files when all tests were run (not a single file re-run)
    // or if previous results are preserved by "cleanOnRerun: false"
    if (this.options.include != null && (allTestsRun || !this.options.cleanOnRerun)) {
      const coveredFiles = coverageMap.files()
      const uncoveredCoverage = await this.getCoverageMapForUncoveredFiles(coveredFiles)

      coverageMap.merge(await transformCoverage(uncoveredCoverage))
    }

    coverageMap.filter((filename) => {
      const exists = existsSync(filename)

      if (this.options.excludeAfterRemap) {
        return exists && this.isIncluded(filename)
      }

      return exists
    })

    if (debug.enabled) {
      debug('Generate coverage total time %d ms', (performance.now() - start!).toFixed())
    }

    return coverageMap
  }

  async generateReports(coverageMap: CoverageMap, allTestsRun: boolean | undefined): Promise<void> {
    const context = libReport.createContext({
      dir: this.options.reportsDirectory,
      coverageMap,
      watermarks: this.options.watermarks,
    })

    if (this.hasTerminalReporter(this.options.reporter)) {
      this.ctx.logger.log(c.blue(' % ') + c.dim('Coverage report from ') + c.yellow(this.name))
    }

    for (const reporter of this.options.reporter) {
      // Type assertion required for custom reporters
      const reportInstance = await libReport.createAsync(
        reporter[0] as Parameters<typeof libReport.create>[0],
        {
          skipFull: this.options.skipFull,
          projectRoot: this.ctx.config.root,
          ...reporter[1],
        },
      )

      reportInstance.execute(context)
    }

    if (this.options.thresholds) {
      await this.reportThresholds(coverageMap, allTestsRun)
    }
  }

  async parseConfigModule(configFilePath: string): Promise<ProxifiedModule<any>> {
    const contents = await fs.readFile(configFilePath, 'utf8')

    return parseModule(`${contents}${this.autoUpdateMarker}`)
  }

  private async getCoverageMapForUncoveredFiles(coveredFiles: string[]) {
    const uncoveredFiles = await this.getUntestedFiles(coveredFiles)

    const cacheKey = Date.now()
    const coverageMap = this.createCoverageMap()

    const transform = this.createUncoveredFileTransformer(this.ctx)

    // Note that these cannot be run parallel as synchronous instrumenter.lastFileCoverage
    // returns the coverage of the last transformed file
    for (const [index, filename] of uncoveredFiles.entries()) {
      let timeout: ReturnType<typeof setTimeout> | undefined
      let start: number | undefined

      if (debug.enabled) {
        start = performance.now()
        timeout = setTimeout(
          () => debug(c.bgRed(`File "${filename}" is taking longer than 3s`)),
          3_000,
        )

        debug('Uncovered file %d/%d', index, uncoveredFiles.length)
      }

      // Make sure file is not served from cache so that instrumenter loads up requested file coverage
      await transform(`${filename}?cache=${cacheKey}&vitest-uncovered-coverage=true`)
      const lastCoverage = this.instrumenter.lastFileCoverage()
      coverageMap.addFileCoverage(lastCoverage)

      if (debug.enabled) {
        clearTimeout(timeout)

        const diff = performance.now() - start!
        const color = diff > 500 ? c.bgRed : c.bgGreen
        debug(`${color(` ${diff.toFixed()} ms `)} ${filename}`)
      }
    }

    return coverageMap
  }

  // the coverage can be enabled after the tests are run
  // this means the coverage will not be injected because the modules are cached,
  // so we are invalidating all modules that don't have the istanbul coverage injected
  onEnabled(): void {
    const environments = this.ctx.projects.flatMap((project) => [
      ...Object.values(project.vite.environments),
      ...Object.values(project.browser?.vite.environments || {}),
    ])

    const seen = new Set<Vite.EnvironmentModuleNode>()
    environments.forEach((environment) => {
      environment.moduleGraph.idToModuleMap.forEach((node) => {
        this.invalidateTree(node, environment.moduleGraph, seen)
      })
    })
  }

  private invalidateTree(
    node: Vite.EnvironmentModuleNode,
    moduleGraph: Vite.EnvironmentModuleGraph,
    seen: Set<Vite.EnvironmentModuleNode>,
  ) {
    if (seen.has(node)) {
      return
    }
    if (node.id && !this.transformedModuleIds.has(node.id)) {
      moduleGraph.invalidateModule(node, seen)
    }
    seen.add(node) // to avoid infinite loops in circular dependencies
    node.importedModules.forEach((mod) => {
      this.invalidateTree(mod, moduleGraph, seen)
    })
  }
}

async function transformCoverage(coverageMap: CoverageMap) {
  const sourceMapStore = libSourceMaps.createSourceMapStore()
  return await sourceMapStore.transformCoverage(coverageMap)
}

/**
 * Remove possible query parameters from filenames
 * - From `/src/components/Header.component.ts?vue&type=script&src=true&lang.ts`
 * - To `/src/components/Header.component.ts`
 */
function removeQueryParameters(filename: string) {
  return filename.split('?')[0]
}