// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
//
// SPDX-License-Identifier: Apache-2.0
import type { AudioFrame } from '@livekit/rtc-node';
import { ThrowsPromise } from '@livekit/throws-transformer/throws';
import type { TypedEventEmitter as TypedEmitter } from '@livekit/typed-emitter';
import type { Span } from '@opentelemetry/api';
import { EventEmitter } from 'node:events';
import type { ReadableStream } from 'node:stream/web';
import { APIConnectionError, APIError } from '../_exceptions.js';
import { log } from '../log.js';
import type { TTSMetrics } from '../metrics/base.js';
import { DeferredReadableStream } from '../stream/deferred_stream.js';
import { recordException, traceTypes, tracer } from '../telemetry/index.js';
import {
  type APIConnectOptions,
  DEFAULT_API_CONNECT_OPTIONS,
  USERDATA_TTS_STARTED_TIME,
  intervalForRetry,
} from '../types.js';
import { AsyncIterableQueue, delay, mergeFrames, startSoon, toError } from '../utils.js';
import type { TimedString } from '../voice/io.js';

/**
 * SynthesizedAudio is a packet of speech synthesis as returned by the TTS.
 */
export interface SynthesizedAudio {
  /** Request ID (one segment could be made up of multiple requests) */
  requestId: string;
  /** Segment ID, each segment is separated by a flush */
  segmentId: string;
  /** Synthesized audio frame */
  frame: AudioFrame;
  /** Current segment of the synthesized audio */
  deltaText?: string;
  /** Whether this is the last frame of the segment (streaming only) */
  final: boolean;
  /**
   * Timed transcripts associated with this audio packet (word-level timestamps).
   */
  timedTranscripts?: TimedString[];
}

/**
 * Describes the capabilities of the TTS provider.
 *
 * @remarks
 * At present, only `streaming` is supplied to this interface, and the framework only supports
 * providers that do have a streaming endpoint.
 */
export interface TTSCapabilities {
  streaming: boolean;
  // Whether this TTS supports aligned transcripts (word-level timestamps).
  alignedTranscript?: boolean;
}

export interface TTSError {
  type: 'tts_error';
  timestamp: number;
  label: string;
  error: Error;
  recoverable: boolean;
}

export type TTSCallbacks = {
  ['metrics_collected']: (metrics: TTSMetrics) => void;
  ['error']: (error: TTSError) => void;
};

/**
 * Time at which text was first sent to the TTS provider, captured on both
 * clocks used for TTFB accounting: `time` on the `performance.now()` scale
 * (seconds, stamped on audio frames) and `hrTime` from
 * `process.hrtime.bigint()` (used for metrics durations).
 * @internal
 */
export interface SynthesizeStreamStartedTime {
  time: number;
  hrTime: bigint;
}

/**
 * An instance of a text-to-speech adapter.
 *
 * @remarks
 * This class is abstract, and as such cannot be used directly. Instead, use a provider plugin that
 * exports its own child TTS class, which inherits this class's methods.
 */
export abstract class TTS extends (EventEmitter as new () => TypedEmitter<TTSCallbacks>) {
  #capabilities: TTSCapabilities;
  #sampleRate: number;
  #numChannels: number;
  abstract label: string;

  constructor(sampleRate: number, numChannels: number, capabilities: TTSCapabilities) {
    super();
    this.#capabilities = capabilities;
    this.#sampleRate = sampleRate;
    this.#numChannels = numChannels;
  }

  /** Returns this TTS's capabilities */
  get capabilities(): TTSCapabilities {
    return this.#capabilities;
  }

  /** Returns the sample rate of audio frames returned by this TTS */
  get sampleRate(): number {
    return this.#sampleRate;
  }

  /** Returns the channel count of audio frames returned by this TTS */
  get numChannels(): number {
    return this.#numChannels;
  }

  /**
   * Get the model name/identifier for this TTS instance.
   *
   * @returns The model name if available, "unknown" otherwise.
   *
   * @remarks
   * Plugins should override this property to provide their model information.
   */
  get model(): string {
    return 'unknown';
  }

  /**
   * Get the provider name for this TTS instance.
   *
   * @returns The provider name if available, "unknown" otherwise.
   *
   * @remarks
   * Plugins should override this property to provide their provider information.
   */
  get provider(): string {
    return 'unknown';
  }

  /**
   * Receives text and returns synthesis in the form of a {@link ChunkedStream}
   */
  abstract synthesize(
    text: string,
    connOptions?: APIConnectOptions,
    abortSignal?: AbortSignal,
  ): ChunkedStream;

  /**
   * Returns a {@link SynthesizeStream} that can be used to push text and receive audio data
   *
   * @param options - Optional configuration including connection options
   */
  abstract stream(options?: { connOptions?: APIConnectOptions }): SynthesizeStream;

  async close(): Promise<void> {
    return;
  }
}

/**
 * An instance of a text-to-speech stream, as an asynchronous iterable iterator.
 *
 * @example Looping through frames
 * ```ts
 * for await (const event of stream) {
 *   await source.captureFrame(event.frame);
 * }
 * ```
 *
 * @remarks
 * This class is abstract, and as such cannot be used directly. Instead, use a provider plugin that
 * exports its own child SynthesizeStream class, which inherits this class's methods.
 */
export abstract class SynthesizeStream
  implements AsyncIterableIterator<SynthesizedAudio | typeof SynthesizeStream.END_OF_STREAM>
{
  protected static readonly FLUSH_SENTINEL = Symbol('FLUSH_SENTINEL');
  static readonly END_OF_STREAM = Symbol('END_OF_STREAM');
  protected input = new AsyncIterableQueue<string | typeof SynthesizeStream.FLUSH_SENTINEL>();
  protected queue = new AsyncIterableQueue<
    SynthesizedAudio | typeof SynthesizeStream.END_OF_STREAM
  >();
  protected output = new AsyncIterableQueue<
    SynthesizedAudio | typeof SynthesizeStream.END_OF_STREAM
  >();
  protected closed = false;
  protected connOptions: APIConnectOptions;
  protected abortController = new AbortController();

  private deferredInputStream: DeferredReadableStream<
    string | typeof SynthesizeStream.FLUSH_SENTINEL
  >;
  private logger = log();

  abstract label: string;

  #tts: TTS;
  #metricsPendingTexts: string[] = [];
  #metricsText = '';
  #monitorMetricsTask?: Promise<void>;
  #ttsRequestSpan?: Span;
  #inputTokens = 0;
  #outputTokens = 0;
  // Deduped provider-known segment ids for the current attempt.
  #providerRequestIds: string[] = [];
  // Current `tts_request_run` span – noteProviderRequestId() updates the
  // attribute on this span as new provider ids become known.
  #currentAttemptSpan?: Span;
  #startedTime?: number;
  #startedHrTime?: bigint;

  constructor(tts: TTS, connOptions: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS) {
    this.#tts = tts;
    this.connOptions = connOptions;
    this.deferredInputStream = new DeferredReadableStream();
    this.pumpInput();

    this.abortController.signal.addEventListener('abort', () => {
      this.deferredInputStream.detachSource();
      // TODO (AJS-36) clean this up when we refactor with streams
      if (!this.input.closed) this.input.close();
      if (!this.output.closed) this.output.close();
      this.closed = true;
    });

    // this is a hack to immitate asyncio.create_task so that mainTask
    // is run **after** the constructor has finished. Otherwise we get
    // runtime error when trying to access class variables in the
    // `run` method.
    // Ensure `this.output` is closed once mainTask settles, even when
    // `monitorMetrics` never started (no `pushText` was ever called).
    // Without this, consumers iterating the stream hang forever.
    startSoon(async () => {
      try {
        await this.mainTask();
      } catch {
        // already surfaced via emitError; swallow to avoid unhandled rejection.
      } finally {
        this.queue.close();
        if (this.#monitorMetricsTask) {
          await this.#monitorMetricsTask.catch(() => {});
        }
        if (!this.output.closed) {
          this.output.close();
        }
      }
    });
  }

  private _mainTaskImpl = async (span: Span) => {
    this.#ttsRequestSpan = span;
    span.setAttributes({
      [traceTypes.ATTR_TTS_STREAMING]: true,
      [traceTypes.ATTR_TTS_LABEL]: this.#tts.label,
    });

    for (let i = 0; i < this.connOptions.maxRetry + 1; i++) {
      try {
        return await tracer.startActiveSpan(
          async (attemptSpan) => {
            attemptSpan.setAttribute(traceTypes.ATTR_RETRY_COUNT, i);
            this.#providerRequestIds = [];
            this.#currentAttemptSpan = attemptSpan;
            try {
              return await this.run();
            } catch (error) {
              recordException(attemptSpan, toError(error));
              throw error;
            } finally {
              this.#currentAttemptSpan = undefined;
            }
          },
          { name: 'tts_request_run' },
        );
      } catch (error) {
        if (error instanceof APIError) {
          const retryInterval = intervalForRetry(this.connOptions, i);

          if (this.connOptions.maxRetry === 0 || !error.retryable) {
            this.emitError({ error, recoverable: false });
            throw error;
          } else if (i === this.connOptions.maxRetry) {
            this.emitError({ error, recoverable: false });
            throw new APIConnectionError({
              message: `failed to generate TTS completion after ${this.connOptions.maxRetry + 1} attempts`,
              options: { retryable: false },
            });
          } else {
            // Don't emit error event for recoverable errors during retry loop
            // to avoid ERR_UNHANDLED_ERROR or premature session termination
            this.logger.warn(
              { tts: this.#tts.label, attempt: i + 1, error },
              `failed to synthesize speech, retrying in ${retryInterval}ms`,
            );
          }

          if (retryInterval > 0) {
            await delay(retryInterval);
          }
        } else {
          this.emitError({ error: toError(error), recoverable: false });
          throw error;
        }
      }
    }
  };

  private mainTask = async () =>
    tracer.startActiveSpan(async (span) => this._mainTaskImpl(span), {
      name: 'tts_request',
      endOnExit: false,
    });

  private emitError({ error, recoverable }: { error: Error; recoverable: boolean }) {
    this.#tts.emit('error', {
      type: 'tts_error',
      timestamp: Date.now(),
      label: this.#tts.label,
      error,
      recoverable,
    });
  }

  /**
   * Record a provider-known id for this stream on the current `tts_request_run`
   * span.
   *
   * Exposed on the span as `lk.provider_request_ids` so users can correlate
   * traces with the provider's server-side logs for debugging. Plugins should
   * call this when the provider-known id becomes available – either the
   * outbound `context_id`/`segment_id` (echoed back by the provider) or a
   * `request_id` / `session_id` returned in a response message. Deduped
   * internally, so calling on every message is fine.
   */
  protected noteProviderRequestId(contextId: string): void {
    if (!contextId || this.#providerRequestIds.includes(contextId)) {
      return;
    }
    this.#providerRequestIds.push(contextId);
    if (this.#currentAttemptSpan && this.#currentAttemptSpan.isRecording()) {
      this.#currentAttemptSpan.setAttribute(
        traceTypes.ATTR_PROVIDER_REQUEST_IDS,
        this.#providerRequestIds,
      );
    }
  }

  /**
   * Mark the time at which text was first sent to the TTS provider for the
   * current segment. Used as the anchor for TTFB so upstream latency (LLM
   * streaming, sentence tokenization) is not attributed to the TTS. Only the
   * first call takes effect; the mark is reset after metrics are emitted for
   * the segment.
   *
   * Plugins should call this right before sending text to the provider.
   * Adapters wrapping another stream may pass that stream's
   * {@link SynthesizeStream.startedTime} to anchor on the wrapped provider.
   */
  protected markStarted(startedTime?: SynthesizeStreamStartedTime): void {
    if (this.#startedTime === undefined) {
      this.#startedTime = startedTime?.time ?? performance.now() / 1000;
      this.#startedHrTime = startedTime?.hrTime ?? process.hrtime.bigint();
    }
  }

  /**
   * Time at which text was first sent to the provider for the current
   * segment, or `undefined` if nothing has been sent (or metrics for the
   * segment were already emitted).
   * @internal
   */
  get startedTime(): SynthesizeStreamStartedTime | undefined {
    if (this.#startedTime === undefined || this.#startedHrTime === undefined) {
      return undefined;
    }
    return { time: this.#startedTime, hrTime: this.#startedHrTime };
  }

  // NOTE(AJS-37): The implementation below uses an AsyncIterableQueue (`this.input`)
  // bridged from a DeferredReadableStream (`this.deferredInputStream`) rather than
  // consuming the stream directly.
  //
  // A full refactor to native Web Streams was considered but is currently deferred.
  //
  // For more context, see the discussion in GitHub issue # 844.
  protected async pumpInput() {
    const reader = this.deferredInputStream.stream.getReader();
    try {
      while (true) {
        const { done, value } = await reader.read();
        if (done || value === SynthesizeStream.FLUSH_SENTINEL) {
          break;
        }
        this.pushText(value);
      }
      this.endInput();
    } catch (error) {
      this.logger.error(error, 'Error reading deferred input stream');
    } finally {
      reader.releaseLock();
      // Ensure output is closed when the stream ends
      if (!this.#monitorMetricsTask) {
        // No text was received, close the output directly
        this.output.close();
      }
    }
  }

  /**
   * Set token usage for token-based TTS billing (e.g., OpenAI TTS).
   * Plugins should call this method to report token usage.
   */
  protected setTokenUsage({
    inputTokens = 0,
    outputTokens = 0,
  }: { inputTokens?: number; outputTokens?: number } = {}): void {
    this.#inputTokens = inputTokens;
    this.#outputTokens = outputTokens;
  }

  protected async monitorMetrics() {
    let audioDurationMs = 0;
    let ttfb: bigint = BigInt(-1);
    let requestId = '';

    const emit = () => {
      if (this.#startedHrTime === undefined) {
        return;
      }
      if (this.#metricsPendingTexts.length) {
        const text = this.#metricsPendingTexts.shift()!;
        const duration = process.hrtime.bigint() - this.#startedHrTime;
        const roundedAudioDurationMs = Math.round(audioDurationMs);
        const metrics: TTSMetrics = {
          type: 'tts_metrics',
          timestamp: Date.now(),
          requestId,
          ttfbMs: ttfb === BigInt(-1) ? -1 : Math.trunc(Number(ttfb / BigInt(1000000))),
          durationMs: Math.trunc(Number(duration / BigInt(1000000))),
          charactersCount: text.length,
          audioDurationMs: roundedAudioDurationMs,
          cancelled: this.abortController.signal.aborted,
          label: this.#tts.label,
          inputTokens: this.#inputTokens,
          outputTokens: this.#outputTokens,
          streamed: true,
          metadata: {
            modelProvider: this.#tts.provider,
            modelName: this.#tts.model,
          },
        };
        if (this.#ttsRequestSpan) {
          this.#ttsRequestSpan.setAttribute(traceTypes.ATTR_TTS_METRICS, JSON.stringify(metrics));
        }
        this.#tts.emit('metrics_collected', metrics);

        // Reset token usage after emitting metrics for the next segment
        this.#inputTokens = 0;
        this.#outputTokens = 0;
        this.#startedTime = undefined;
        this.#startedHrTime = undefined;
        ttfb = BigInt(-1);
        audioDurationMs = 0;
      }
    };

    for await (const audio of this.queue) {
      if (this.abortController.signal.aborted) {
        break;
      }
      if (audio === SynthesizeStream.END_OF_STREAM) {
        this.output.put(audio);
        continue;
      }
      if (this.#startedTime !== undefined) {
        audio.frame.userdata[USERDATA_TTS_STARTED_TIME] = this.#startedTime;
      }
      this.output.put(audio);
      requestId = audio.requestId;
      // The Python AudioEmitter records each `segment_id` automatically via
      // `start_segment`. JS plugins emit `segmentId` directly on the audio
      // frame, so we mirror that behavior here.
      if (audio.segmentId) {
        this.noteProviderRequestId(audio.segmentId);
      }
      if (ttfb === BigInt(-1) && this.#startedHrTime !== undefined) {
        ttfb = process.hrtime.bigint() - this.#startedHrTime;
      }
      // TODO(AJS-102): use frame.durationMs once available in rtc-node
      audioDurationMs += (audio.frame.samplesPerChannel / audio.frame.sampleRate) * 1000;
      if (audio.final) {
        emit();
      }
    }

    if (requestId) {
      emit();
    }

    if (this.#ttsRequestSpan) {
      this.#ttsRequestSpan.end();
      this.#ttsRequestSpan = undefined;
    }
  }

  protected abstract run(): Promise<void>;

  updateInputStream(text: ReadableStream<string>) {
    this.deferredInputStream.setSource(text);
  }

  /** Push a string of text to the TTS */
  /** @deprecated Use `updateInputStream` instead */
  pushText(text: string) {
    if (!this.#monitorMetricsTask) {
      this.#monitorMetricsTask = this.monitorMetrics();
      // Close output when metrics task completes
      this.#monitorMetricsTask.finally(() => this.output.close());
    }
    this.#metricsText += text;

    if (this.input.closed || this.closed) {
      // Stream was aborted/closed, silently skip
      return;
    }

    this.input.put(text);
  }

  /** Flush the TTS, causing it to process all pending text */
  flush() {
    if (this.#metricsText) {
      this.#metricsPendingTexts.push(this.#metricsText);
      this.#metricsText = '';
    }

    if (this.input.closed || this.closed) {
      // Stream was aborted/closed, silently skip
      return;
    }

    this.input.put(SynthesizeStream.FLUSH_SENTINEL);
  }

  /** Mark the input as ended and forbid additional pushes */
  endInput() {
    this.flush();

    if (this.input.closed || this.closed) {
      // Stream was aborted/closed, silently skip
      return;
    }

    this.input.close();
  }

  next(): Promise<IteratorResult<SynthesizedAudio | typeof SynthesizeStream.END_OF_STREAM>> {
    return this.output.next();
  }

  get abortSignal(): AbortSignal {
    return this.abortController.signal;
  }

  /** Close both the input and output of the TTS stream */
  close() {
    this.abortController.abort();
  }

  [Symbol.asyncIterator](): SynthesizeStream {
    return this;
  }
}

/**
 * An instance of a text-to-speech response, as an asynchronous iterable iterator.
 *
 * @example Looping through frames
 * ```ts
 * for await (const event of stream) {
 *   await source.captureFrame(event.frame);
 * }
 * ```
 *
 * @remarks
 * This class is abstract, and as such cannot be used directly. Instead, use a provider plugin that
 * exports its own child ChunkedStream class, which inherits this class's methods.
 */
export abstract class ChunkedStream implements AsyncIterableIterator<SynthesizedAudio> {
  protected queue = new AsyncIterableQueue<SynthesizedAudio>();
  protected output = new AsyncIterableQueue<SynthesizedAudio>();
  protected closed = false;
  abstract label: string;
  #text: string;
  #tts: TTS;
  #ttsRequestSpan?: Span;
  private _connOptions: APIConnectOptions;
  private logger = log();
  #inputTokens = 0;
  #outputTokens = 0;
  #startedTime: number;

  protected abortController = new AbortController();

  constructor(
    text: string,
    tts: TTS,
    connOptions: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
    abortSignal?: AbortSignal,
  ) {
    this.#text = text;
    this.#tts = tts;
    this._connOptions = connOptions;
    this.#startedTime = performance.now() / 1000;

    if (abortSignal) {
      abortSignal.addEventListener('abort', () => this.abortController.abort(), { once: true });
    }

    this.monitorMetrics();

    // this is a hack to immitate asyncio.create_task so that mainTask
    // is run **after** the constructor has finished. Otherwise we get
    // runtime error when trying to access class variables in the
    // `run` method.
    ThrowsPromise.resolve().then(() => this.mainTask().finally(() => this.queue.close()));
  }

  private _mainTaskImpl = async (span: Span) => {
    this.#ttsRequestSpan = span;
    span.setAttributes({
      [traceTypes.ATTR_TTS_STREAMING]: false,
      [traceTypes.ATTR_TTS_LABEL]: this.#tts.label,
    });

    for (let i = 0; i < this._connOptions.maxRetry + 1; i++) {
      try {
        return await tracer.startActiveSpan(
          async (attemptSpan) => {
            attemptSpan.setAttribute(traceTypes.ATTR_RETRY_COUNT, i);
            try {
              return await this.run();
            } catch (error) {
              recordException(attemptSpan, toError(error));
              throw error;
            }
          },
          { name: 'tts_request_run' },
        );
      } catch (error) {
        if (error instanceof APIError) {
          const retryInterval = intervalForRetry(this._connOptions, i);

          if (this._connOptions.maxRetry === 0 || !error.retryable) {
            this.emitError({ error, recoverable: false });
            throw error;
          } else if (i === this._connOptions.maxRetry) {
            this.emitError({ error, recoverable: false });
            throw new APIConnectionError({
              message: `failed to generate TTS completion after ${this._connOptions.maxRetry + 1} attempts`,
              options: { retryable: false },
            });
          } else {
            // Don't emit error event for recoverable errors during retry loop
            // to avoid ERR_UNHANDLED_ERROR or premature session termination
            this.logger.warn(
              { tts: this.#tts.label, attempt: i + 1, error },
              `failed to generate TTS completion, retrying in ${retryInterval}ms`,
            );
          }

          if (retryInterval > 0) {
            await delay(retryInterval);
          }
        } else {
          this.emitError({ error: toError(error), recoverable: false });
          throw error;
        }
      }
    }
  };

  private async mainTask() {
    return tracer.startActiveSpan(async (span) => this._mainTaskImpl(span), {
      name: 'tts_request',
      endOnExit: false,
    });
  }

  private emitError({ error, recoverable }: { error: Error; recoverable: boolean }) {
    this.#tts.emit('error', {
      type: 'tts_error',
      timestamp: Date.now(),
      label: this.#tts.label,
      error,
      recoverable,
    });
  }

  protected abstract run(): Promise<void>;

  get inputText(): string {
    return this.#text;
  }

  get abortSignal(): AbortSignal {
    return this.abortController.signal;
  }

  /**
   * Set token usage for token-based TTS billing (e.g., OpenAI TTS).
   * Plugins should call this method to report token usage.
   */
  protected setTokenUsage({
    inputTokens = 0,
    outputTokens = 0,
  }: { inputTokens?: number; outputTokens?: number } = {}): void {
    this.#inputTokens = inputTokens;
    this.#outputTokens = outputTokens;
  }

  protected async monitorMetrics() {
    const startTime = process.hrtime.bigint();
    let audioDurationMs = 0;
    let ttfb: bigint = BigInt(-1);
    let requestId = '';

    for await (const audio of this.queue) {
      audio.frame.userdata[USERDATA_TTS_STARTED_TIME] = this.#startedTime;
      this.output.put(audio);
      requestId = audio.requestId;
      if (ttfb === BigInt(-1)) {
        ttfb = process.hrtime.bigint() - startTime;
      }
      audioDurationMs += (audio.frame.samplesPerChannel / audio.frame.sampleRate) * 1000;
    }
    this.output.close();

    const duration = process.hrtime.bigint() - startTime;
    const metrics: TTSMetrics = {
      type: 'tts_metrics',
      timestamp: Date.now(),
      requestId,
      ttfbMs: ttfb === BigInt(-1) ? -1 : Math.trunc(Number(ttfb / BigInt(1000000))),
      durationMs: Math.trunc(Number(duration / BigInt(1000000))),
      charactersCount: this.#text.length,
      audioDurationMs: Math.round(audioDurationMs),
      cancelled: false, // TODO(AJS-186): support ChunkedStream with 1.0 - add this.abortController.signal.aborted here
      label: this.#tts.label,
      inputTokens: this.#inputTokens,
      outputTokens: this.#outputTokens,
      streamed: false,
      metadata: {
        modelProvider: this.#tts.provider,
        modelName: this.#tts.model,
      },
    };

    if (this.#ttsRequestSpan) {
      this.#ttsRequestSpan.setAttribute(traceTypes.ATTR_TTS_METRICS, JSON.stringify(metrics));
      this.#ttsRequestSpan.end();
      this.#ttsRequestSpan = undefined;
    }

    this.#tts.emit('metrics_collected', metrics);
  }

  /** Collect every frame into one in a single call */
  async collect(): Promise<AudioFrame> {
    const frames = [];
    for await (const event of this) {
      frames.push(event.frame);
    }
    return mergeFrames(frames);
  }

  next(): Promise<IteratorResult<SynthesizedAudio>> {
    return this.output.next();
  }

  /** Close both the input and output of the TTS stream */
  close() {
    if (!this.queue.closed) this.queue.close();
    if (!this.output.closed) this.output.close();
    if (!this.abortController.signal.aborted) this.abortController.abort();
    this.closed = true;
  }

  [Symbol.asyncIterator](): ChunkedStream {
    return this;
  }
}
