import { CompletionOptions, ILLM } from "../.."
import { StreamTransformPipeline } from "../filtering/streamTransforms/StreamTransformPipeline"
import { HelperVars } from "../util/HelperVars"

import { GeneratorReuseManager } from "./GeneratorReuseManager"
import { stopAfterMaxProcessingTime } from "./utils"

export class CompletionStreamer {
  private streamTransformPipeline = new StreamTransformPipeline()
  private generatorReuseManager: GeneratorReuseManager

  constructor(onError: (err: unknown) => void) {
    this.generatorReuseManager = new GeneratorReuseManager(onError)
  }

  async *streamCompletionWithFilters(
    token: AbortSignal,
    llm: ILLM,
    prefix: string,
    suffix: string,
    prompt: string,
    multiline: boolean,
    completionOptions: Partial<CompletionOptions> | undefined,
    helper: HelperVars,
  ) {
    // Full stop means to stop the LLM's generation, instead of just truncating the displayed completion
    const fullStop = () => this.generatorReuseManager.currentGenerator?.cancel()

    // Try to reuse pending requests if what the user typed matches start of completion
    const generator = this.generatorReuseManager.getGenerator(
      prefix,
      (abortSignal: AbortSignal) => {
        const generator = llm.supportsFim()
          ? llm.streamFim(prefix, suffix, abortSignal, completionOptions)
          : llm.streamComplete(prompt, abortSignal, {
              ...completionOptions,
              raw: true,
            })

        /**
         * This transformer applies even on reused generator. We are deliberately
         * not using streamTransformPipeline because we want to capture and stop
         * the request even if the generator is being reused.
         */
        return helper.options.transform
          ? stopAfterMaxProcessingTime(generator, helper.options.modelTimeout * 2.5, fullStop)
          : generator
      },
      multiline,
    )

    // LLM
    const generatorWithCancellation = async function* () {
      for await (const update of generator) {
        if (token.aborted) {
          return
        }
        yield update
      }
    }

    const initialGenerator = generatorWithCancellation()
    const transformedGenerator = helper.options.transform
      ? this.streamTransformPipeline.transform(
          initialGenerator,
          prefix,
          suffix,
          multiline,
          completionOptions?.stop || [],
          fullStop,
          helper,
        )
      : initialGenerator

    for await (const update of transformedGenerator) {
      yield update
    }
  }
}
