'use strict'

const path = require('bare-path')
const BaseInference = require('@qvac/infer-base/WeightsProvider/BaseInference')
const WeightsProvider = require('@qvac/infer-base/WeightsProvider/WeightsProvider')

const { TranslationInterface } = require('./marian')
const QvacResponse = require('@qvac/response')
const { IndicProcessor } = require('./third-party/indic-processor')

const END_OF_INPUT = 'end of job'

class QvacIndicTransResponse extends QvacResponse {
  /**
     * Creates an instance of QvacIndicTransResponse.
     * @constructor
     * @param {IndicProcessor} processor
     * @param {string} dstLang
     * @param {Object} handlers
     */
  constructor (processor, dstLang, handlers) {
    super(handlers)
    this.processor = processor
    this.dstLang = dstLang
  }

  onCancel (callback) {
    return super.onCancel(callback)
  }

  onError (callback) {
    return super.onError(callback)
  }

  onFinish (callback) {
    return super.onFinish(callback)
  }

  onUpdate (callback) {
    return super.onUpdate((data) => {
      const [postProcessedText] = this.processor.postprocessBatch(
        [data],
        this.dstLang
      )
      return callback(postProcessedText)
    })
  }

  async * iterate () {
    for await (const output of super.iterate()) {
      const [postProcessedText] = this.processor.postprocessBatch(
        [output],
        this.dstLang
      )
      yield postProcessedText
    }
  }
}

/**
 * TranslationNmtcpp implementation for Marian/IndicTrans translation model
 */
class TranslationNmtcpp extends BaseInference {
  /**
   * Available model types for translation
   * @static
   * @type {Object}
   * @property {string} IndicTrans
   * @property {string} Opus
   */
  static ModelTypes = {
    IndicTrans: 'IndicTrans',
    Opus: 'Opus'
  }

  /**
     * Creates an instance of MLCInference.
     * @constructor
   *
     * @param {Object} args arguments for inference setup
     * @param {Object} config - environment specific inference setup configuration
     */
  constructor ({ loader, diskPath, modelName, params, logger = null, exclusiveRun = true, ...args }, config = {}) {
    super({ logger, exclusiveRun, ...args })
    this.loader = loader
    this.weightsProvider = new WeightsProvider(loader, this.logger)
    const { modelType, ...additionalConfig } = config
    this._modelType = modelType
    this._config = additionalConfig
    this._diskPath = diskPath
    this._modelName = modelName
    this._params = params
  }

  async _load (close = false, reportProgressCallback) {
    await this.loader.ready()
    try {
      await this.downloadWeights(reportProgressCallback)

      const configurationParams = {
        path: this._config.path || path.join(this._diskPath, this._modelName),
        config: this._config
      }
      this.addon = this.createAddon(configurationParams)
      await this.addon.activate()
      this.state.configLoaded = true
    } finally {
      if (close) {
        await this.loader.close()
      }
    }
  }

  async _runInternal (input) {
    if (this._modelType === TranslationNmtcpp.ModelTypes.IndicTrans) {
      const processor = new IndicProcessor()
      const [processedText] = processor.preprocessBatch(
        [input],
        this._params.srcLang,
        this._params.dstLang
      )
      const jobId = await this.addon.append({
        type: 'text',
        input: processedText
      })
      const response = new QvacIndicTransResponse(
        processor,
        this._params.dstLang,
        {
          cancelHandler: () => {
            return this.addon.cancel(jobId)
          },
          pauseHandler: () => {
            return this.addon.pause()
          },
          continueHandler: () => {
            return this.addon.activate()
          }
        }
      )

      this._saveJobToResponseMapping(jobId, response)
      await this.addon.append({ type: END_OF_INPUT })
      return response
    }

    const jobId = await this.addon.append({ type: 'text', input })

    const response = new QvacResponse({
      cancelHandler: () => {
        return this.addon.cancel(jobId)
      },
      pauseHandler: () => {
        return this.addon.pause()
      },
      continueHandler: () => {
        return this.addon.activate()
      }
    })

    this._saveJobToResponseMapping(jobId, response)
    await this.addon.append({ type: END_OF_INPUT })
    return response
  }

  createAddon (configurationParams) {
    return new TranslationInterface(
      configurationParams,
      this._outputCallback.bind(this),
      this.logger
    )
  }

  async _downloadWeights (reportProgressCallback) {
    const models = [this._modelName]

    this.logger.info('Loading weight files:', models)

    const result = await this.weightsProvider.downloadFiles(models, this._diskPath, {
      closeLoader: true,
      onDownloadProgress: reportProgressCallback
    })
    this.logger.info('Weight files downloaded successfully', { models })
    return result
  }
}

module.exports = TranslationNmtcpp
