import { LayersModel } from '@tensorflow/tfjs-node'

import { NlpTaskConfig } from '../storage'
import { Locale } from '../types'
import { ModelInfo } from './model-info'

export abstract class ModelSelector<R, C extends NlpTaskConfig> {
  private models: { [locale: string]: R } = {}

  protected constructor(
    readonly locales: Locale[],
    readonly modelsBaseUrl: string
  ) {}

  protected async load(): Promise<this> {
    const modelsInfo = this.loadModelsInfo()
    await this.loadModels(modelsInfo)
    return this
  }

  private loadModelsInfo(): ModelInfo<C>[] {
    return this.locales.map(
      locale => new ModelInfo<C>(locale, `${this.modelsBaseUrl}/${locale}`)
    )
  }

  private async loadModels(modelsInfo: ModelInfo<C>[]): Promise<void> {
    for (const modelInfo of modelsInfo) {
      const locale = modelInfo.locale
      const config = await modelInfo.getConfig()
      const model = await modelInfo.getModel()
      this.models[locale] = this.createRecognizer(config, model)
    }
  }

  select(locale: Locale): R {
    return this.models[locale]
  }

  protected abstract createRecognizer(config: C, model: LayersModel): R
}
