iterative/vscode-dvc

import { EventEmitter } from 'vscode'
import isEmpty from 'lodash.isempty'
import { Experiments } from '.'
import {
  getBranchExperimentCommand,
  getPushExperimentCommand
} from './commands'
import { TableData } from './webview/contract'
import { Args } from '../cli/dvc/constants'
import { AvailableCommands, CommandId } from '../commands/internal'
import { ResourceLocator } from '../resourceLocator'
import { Setup } from '../setup'
import { Toast } from '../vscode/toast'
import { getInput, getPositiveIntegerInput } from '../vscode/inputBox'
import { BaseWorkspaceWebviews } from '../webview/workspace'
import { Title } from '../vscode/title'
import { quickPickManyValues } from '../vscode/quickPick'
import { WorkspacePipeline } from '../pipeline/workspace'

export class WorkspaceExperiments extends BaseWorkspaceWebviews<
  Experiments,
  TableData
> {
  public readonly experimentsChanged = this.dispose.track(
    new EventEmitter<void>()
  )

  public readonly onDidChangeExperiments = this.experimentsChanged.event

  public readonly columnsChanged = this.dispose.track(new EventEmitter<void>())
  public readonly columnsOrderOrStatusChanged = this.dispose.track(
    new EventEmitter<void>()
  )

  private focusedFileDvcRoot: string | undefined

  public addFilter(overrideRoot?: string) {
    return this.getRepositoryThenUpdate('addFilter', overrideRoot)
  }

  public addStarredFilter(overrideRoot?: string) {
    return this.getRepositoryThenUpdate('addStarredFilter', overrideRoot)
  }

  public removeFilters() {
    return this.getRepositoryThenUpdate('removeFilters')
  }

  public addSort(overrideRoot?: string) {
    return this.getRepositoryThenUpdate('addSort', overrideRoot)
  }

  public addStarredSort(overrideRoot?: string) {
    return this.getRepositoryThenUpdate('addStarredSort', overrideRoot)
  }

  public removeSorts() {
    return this.getRepositoryThenUpdate('removeSorts')
  }

  public selectExperimentsToPlot(overrideRoot?: string) {
    return this.getRepositoryThenUpdate('selectExperimentsToPlot', overrideRoot)
  }

  public selectColumns(overrideRoot?: string) {
    return this.getRepositoryThenUpdate('selectColumns', overrideRoot)
  }

  public selectFirstColumns(overrideRoot?: string) {
    return this.getRepositoryThenUpdate('selectFirstColumns', overrideRoot)
  }

  public async selectExperimentsToStop() {
    const cwd = await this.getFocusedOrOnlyOrPickProject()
    if (!cwd) {
      return
    }

    const ids = await this.getRepository(cwd).pickRunningExperiments()

    if (!ids || isEmpty(ids)) {
      return
    }
    return this.stopExperiments(cwd, ...ids)
  }

  public stopExperiments(dvcRoot: string, ...ids: string[]) {
    return this.getRepository(dvcRoot).stopExperiments(ids)
  }

  public async selectExperimentsToPush(setup: Setup) {
    const dvcRoot = await this.getFocusedOrOnlyOrPickProject()
    if (!dvcRoot) {
      return
    }

    const ids = await this.getRepository(dvcRoot).pickExperimentsToPush()
    if (!ids || isEmpty(ids)) {
      return
    }

    const pushCommand = getPushExperimentCommand(
      this,
      this.internalCommands,
      setup
    )

    return pushCommand({ dvcRoot, ids })
  }

  public async selectExperimentsToRemove() {
    const cwd = await this.getFocusedOrOnlyOrPickProject()
    if (!cwd) {
      return
    }

    const ids = await this.getRepository(cwd).pickExperimentsToRemove()

    if (!ids || isEmpty(ids)) {
      return
    }
    return this.runCommand(AvailableCommands.EXP_REMOVE, cwd, ...ids)
  }

  public async modifyWorkspaceParamsAndRun(overrideRoot?: string) {
    const project = await this.getDvcRoot(overrideRoot)
    if (!project) {
      return
    }
    const repository = this.getRepository(project)
    if (!repository) {
      return
    }

    return await repository.modifyWorkspaceParamsAndRun()
  }

  public async modifyWorkspaceParamsAndQueue(overrideRoot?: string) {
    const project = await this.getDvcRoot(overrideRoot)
    if (!project) {
      return
    }
    const repository = this.getRepository(project)
    if (!repository) {
      return
    }

    return await repository.modifyWorkspaceParamsAndQueue()
  }

  public async getCwdThenRun(commandId: CommandId) {
    const cwd = await this.getCwd()

    if (!cwd) {
      return 'Could not run task'
    }

    return this.internalCommands.executeCommand(commandId, cwd)
  }

  public getCwdThenReport(commandId: CommandId) {
    return Toast.showOutput(this.getCwdThenRun(commandId))
  }

  public getCwdAndExpNameThenRun(commandId: CommandId) {
    return this.pickExpThenRun(commandId, cwd =>
      this.pickCommitOrExperiment(cwd)
    )
  }

  public async getCwdAndQuickPickThenRun(
    commandId: CommandId,
    quickPick: () => Thenable<string[] | undefined>
  ) {
    const cwd = await this.getCwd()
    if (!cwd) {
      return
    }
    const result = await quickPick()

    if (result) {
      return this.runCommand(commandId, cwd, ...result)
    }
  }

  public async createExperimentBranch() {
    const cwd = await this.getCwd()
    if (!cwd) {
      return
    }

    const experimentId = await this.pickCommitOrExperiment(cwd)

    if (!experimentId) {
      return
    }
    return this.getInputAndRun(
      getBranchExperimentCommand(this),
      Title.ENTER_BRANCH_NAME,
      `${experimentId}-branch`,
      cwd,
      experimentId
    )
  }

  public async getInputAndRun(
    runCommand: (...args: Args) => Promise<void | string>,
    title: Title,
    defaultValue: string,
    ...args: Args
  ) {
    const input = await getInput(title, defaultValue)
    if (!input) {
      return
    }
    return runCommand(...args, input)
  }

  public async getCwdIntegerInputAndRun(
    commandId: CommandId,
    title: Title,
    options: { prompt: string; value: string }
  ) {
    const cwd = await this.getCwd()
    if (!cwd) {
      return
    }

    const integer = await getPositiveIntegerInput(title, options)

    if (!integer) {
      return
    }

    return this.runCommand(commandId, cwd, integer)
  }

  public runCommand(commandId: CommandId, cwd: string, ...args: Args) {
    const output = this.internalCommands.executeCommand(commandId, cwd, ...args)

    void Toast.showOutput(output)
    return output
  }

  public createRepository(
    dvcRoot: string,
    subProjects: string[],
    pipeline: WorkspacePipeline,
    setup: Setup,
    resourceLocator: ResourceLocator
  ) {
    const experiments = this.dispose.track(
      new Experiments(
        dvcRoot,
        this.internalCommands,
        pipeline.getRepository(dvcRoot),
        resourceLocator,
        this.workspaceState,
        (branchesSelected: string[]) => this.selectBranches(branchesSelected),
        subProjects
      )
    )

    this.setRepository(dvcRoot, experiments)

    void experiments.setStudioValues(
      setup.getStudioUrl(),
      setup.getStudioAccessToken()
    )

    experiments.dispose.track(
      setup.onDidChangeStudioConnection(() => {
        void experiments.setStudioValues(
          setup.getStudioUrl(),
          setup.getStudioAccessToken()
        )
      })
    )

    experiments.dispose.track(
      experiments.onDidChangeIsWebviewFocused(
        dvcRoot => (this.focusedWebviewDvcRoot = dvcRoot)
      )
    )

    experiments.dispose.track(
      experiments.onDidChangeIsExperimentsFileFocused(
        dvcRoot => (this.focusedFileDvcRoot = dvcRoot)
      )
    )

    experiments.dispose.track(
      experiments.onDidChangeExperiments(() => {
        this.experimentsChanged.fire()
      })
    )

    experiments.dispose.track(
      experiments.onDidChangeColumns(() => {
        this.columnsChanged.fire()
      })
    )

    experiments.dispose.track(
      experiments.onDidChangeColumnOrderOrStatus(() => {
        this.experimentsChanged.fire()
      })
    )

    return experiments
  }

  public getFocusedOrOnlyOrPickProject() {
    return (
      this.focusedWebviewDvcRoot ||
      this.focusedFileDvcRoot ||
      this.getOnlyOrPickProject()
    )
  }

  public getHasData() {
    const allLoading = undefined

    const repositories = Object.values(this.repositories)

    if (repositories.some(repository => repository.getHasData())) {
      return true
    }

    if (repositories.some(repository => repository.getHasData() === false)) {
      return false
    }

    return allLoading
  }

  public getCliError() {
    const repositories = Object.values(this.repositories)
    const errors = []
    for (const repository of repositories) {
      const cliError = repository.getCliError()
      if (!cliError) {
        continue
      }
      errors.push(cliError)
    }
    return errors.length > 0 ? errors.join('\n') : undefined
  }

  public hasRunningExperiment() {
    return Object.values(this.repositories).some(experiments =>
      experiments.hasRunningExperiment()
    )
  }

  public async selectBranches(branchesSelected: string[]) {
    const dvcRoot = await this.getDvcRoot()
    if (!dvcRoot) {
      return
    }
    const repository = this.getRepository(dvcRoot)

    const allBranches = repository.getAvailableBranchesToSelect()

    return await quickPickManyValues(
      allBranches.map(branch => {
        return {
          label: branch,
          picked: branchesSelected.includes(branch),
          value: branch
        }
      }),
      {
        title: Title.SELECT_BRANCHES
      }
    )
  }

  private async getRepositoryThenUpdate(
    method:
      | 'addFilter'
      | 'addStarredFilter'
      | 'removeFilters'
      | 'addSort'
      | 'addStarredSort'
      | 'removeSorts'
      | 'selectExperimentsToPlot'
      | 'selectColumns'
      | 'selectFirstColumns',
    overrideRoot?: string
  ) {
    const dvcRoot = await this.getDvcRoot(overrideRoot)
    if (!dvcRoot) {
      return
    }
    return this.getRepository(dvcRoot)[method]()
  }

  private async getCwd(overrideRoot?: string) {
    const project = await this.getDvcRoot(overrideRoot)
    if (!project) {
      return
    }
    const repository = this.getRepository(project)
    return await repository.getPipelineCwd()
  }

  private async pickExpThenRun(
    commandId: CommandId,
    pickFunc: (cwd: string) => Thenable<string | undefined> | undefined
  ) {
    const cwd = await this.getFocusedOrOnlyOrPickProject()
    if (!cwd) {
      return
    }

    const experimentId = await pickFunc(cwd)

    if (!experimentId) {
      return
    }
    return this.runCommand(commandId, cwd, experimentId)
  }

  private pickCommitOrExperiment(cwd: string) {
    return this.getRepository(cwd).pickCommitOrExperiment()
  }
}