// SPDX-FileCopyrightText: 2025 LiveKit, Inc.
//
// SPDX-License-Identifier: Apache-2.0
import type { AudioFrame } from '@livekit/rtc-node';
import { ThrowsPromise } from '@livekit/throws-transformer/throws';
import { WebSocket } from 'ws';
import { APIError, APIStatusError, APITimeoutError } from '../_exceptions.js';
import { AudioByteStream } from '../audio.js';
import { ConnectionPool } from '../connection_pool.js';
import { type LanguageCode, normalizeLanguage } from '../language.js';
import { log } from '../log.js';
import { createStreamChannel } from '../stream/stream_channel.js';
import { basic as tokenizeBasic } from '../tokenize/index.js';
import type { ChunkedStream } from '../tts/index.js';
import { SynthesizeStream as BaseSynthesizeStream, TTS as BaseTTS } from '../tts/index.js';
import { type APIConnectOptions, DEFAULT_API_CONNECT_OPTIONS } from '../types.js';
import {
  Event,
  Future,
  Task,
  cancelAndWait,
  combineSignals,
  shortuuid,
  waitUntilTimeout,
} from '../utils.js';
import { type TimedString, createTimedString } from '../voice/io.js';
import {
  type TtsClientEvent,
  type TtsServerEvent,
  ttsClientEventSchema,
  ttsKnownServerEventSchema,
  ttsServerEventSchema,
} from './api_protos.js';
import { type AnyString, connectWs, createAccessToken, getDefaultInferenceUrl } from './utils.js';

export type CartesiaModels =
  | 'cartesia/sonic-3.5'
  | 'cartesia/sonic-3'
  | 'cartesia/sonic-2'
  | 'cartesia/sonic-turbo'
  | 'cartesia/sonic'
  | 'cartesia/sonic-3-latest'
  | 'cartesia/sonic-latest';

export type DeepgramTTSModels = 'deepgram/aura' | 'deepgram/aura-2';

export type ElevenlabsModels =
  | 'elevenlabs/eleven_flash_v2'
  | 'elevenlabs/eleven_flash_v2_5'
  | 'elevenlabs/eleven_turbo_v2'
  | 'elevenlabs/eleven_turbo_v2_5'
  | 'elevenlabs/eleven_multilingual_v2'
  | 'elevenlabs/eleven_v3';

export type InworldModels =
  | 'inworld/inworld-tts-2'
  | 'inworld/inworld-tts-1.5-max'
  | 'inworld/inworld-tts-1.5-mini'
  | 'inworld/inworld-tts-1.5'
  | 'inworld/inworld-tts-1-max'
  | 'inworld/inworld-tts-1';

export type RimeModels = 'rime/arcana' | 'rime/coda' | 'rime/mistv2' | 'rime/mistv3' | 'rime/mist';

export type XaiTTSModels = 'xai/tts-1';

export interface CartesiaOptions {
  emotion?: string;
  /** Maximum duration of audio in seconds. */
  duration?: number;
  /** Speech speed. Default: not specified. */
  speed?: 'slow' | 'normal' | 'fast' | number;
  volume?: number;
  max_buffer_delay_ms?: number;
  add_timestamps?: boolean;
  add_phoneme_timestamps?: boolean;
  use_normalized_timestamps?: boolean;
}

export interface ElevenlabsOptions {
  /** Inactivity timeout in seconds. Default: 60. */
  inactivity_timeout?: number;
  /** Text normalization mode. Default: "auto". */
  apply_text_normalization?: 'auto' | 'off' | 'on';
  auto_mode?: boolean;
  enable_logging?: boolean;
  enable_ssml_parsing?: boolean;
  sync_alignment?: boolean;
  language_code?: string;
  /** Voice stability tuning, typically in the range [0, 1]. */
  stability?: number;
  /** Voice similarity tuning, typically in the range [0, 1]. */
  similarity_boost?: number;
  /** Style exaggeration tuning, typically in the range [0, 1]. */
  style?: number;
  /** Speech speed multiplier. */
  speed?: number;
  use_speaker_boost?: boolean;
  chunk_length_schedule?: number[];
  preferred_alignment?: string;
}

export interface DeepgramTTSOptions {
  /** Default: false. */
  mip_opt_out?: boolean;
}

export interface RimeOptions {
  /** Maximum number of tokens to generate. */
  max_tokens?: number;
  /** Values above 1 slow down audio; values below 1 speed it up. */
  time_scale_factor?: number;
  /** Default 1.0, <1 = faster, >1 = slower. */
  speed_alpha?: number;
  /** Default false. */
  pause_between_brackets?: boolean;
  /** Default false. */
  phonemize_between_brackets?: boolean;
  /** Comma-separated speed factors for [bracketed] words. */
  inline_speed_alpha?: string;
  /** Default false. */
  no_text_normalization?: boolean;
}

export interface InworldOptions {
  /** Range >0.5, <=1.5. */
  speaking_rate?: number;
  /** Range 0-2. */
  temperature?: number;
  /** Controls output variation on inworld-tts-2 only. */
  delivery_mode?: 'DELIVERY_MODE_UNSPECIFIED' | 'STABLE' | 'BALANCED' | 'CREATIVE';
  timestamp_type?: 'TIMESTAMP_TYPE_UNSPECIFIED' | 'WORD' | 'CHARACTER';
  apply_text_normalization?: 'APPLY_TEXT_NORMALIZATION_UNSPECIFIED' | 'ON' | 'OFF';
  /** @deprecated Backward-compatible alias. Use `apply_text_normalization`. */
  text_normalization?: 'ON' | 'OFF';
}

export interface XaiTTSOptions {
  /** Output bit rate in bits per second. */
  bit_rate?: 32000 | 64000 | 96000 | 128000 | 192000;
}

type _TTSModels =
  | CartesiaModels
  | DeepgramTTSModels
  | ElevenlabsModels
  | RimeModels
  | InworldModels
  | XaiTTSModels;

export type TTSModels =
  | CartesiaModels
  | DeepgramTTSModels
  | ElevenlabsModels
  | RimeModels
  | InworldModels
  | XaiTTSModels
  | AnyString;

export type ModelWithVoice = `${_TTSModels}:${string}` | TTSModels;

export type TTSOptions<TModel extends TTSModels> = TModel extends CartesiaModels
  ? CartesiaOptions
  : TModel extends DeepgramTTSModels
    ? DeepgramTTSOptions
    : TModel extends ElevenlabsModels
      ? ElevenlabsOptions
      : TModel extends RimeModels
        ? RimeOptions
        : TModel extends InworldModels
          ? InworldOptions
          : TModel extends XaiTTSModels
            ? XaiTTSOptions
            : Record<string, unknown>;

/** Parse a model string into [model, voice]. Voice is undefined if not specified. */
export function parseTTSModelString(model: string): [string, string | undefined] {
  const idx = model.lastIndexOf(':');
  if (idx !== -1) {
    return [model.slice(0, idx), model.slice(idx + 1)];
  }
  return [model, undefined];
}

/**
 * Whether the given Inference TTS provider+options combination emits aligned
 * (word- or character-level) timestamps alongside the synthesized audio.
 *
 * Mirrors the Python implementation: support is opt-in through the provider's
 * `extra_kwargs` / `modelOptions` payload — the gateway only emits
 * `output_alignment` events when the provider-specific flag is set.
 */
export function hasAlignedTranscript(
  model: string | undefined,
  modelOptions: Record<string, unknown> | undefined,
): boolean {
  if (!model) return false;
  const provider = model.split('/')[0];
  const opts = modelOptions ?? {};
  if (provider === 'cartesia') {
    return Boolean(opts.add_timestamps);
  }
  if (provider === 'elevenlabs') {
    return Boolean(opts.sync_alignment);
  }
  if (provider === 'inworld') {
    return opts.timestamp_type === 'WORD' || opts.timestamp_type === 'CHARACTER';
  }
  return false;
}

/** Inference Fallback Adapter: configuration for a fallback TTS model that runs server-side in LiveKit Inference, providing automatic fallback between providers. */
export interface TTSFallbackModel {
  /** Model name (e.g. "cartesia/sonic", "elevenlabs/eleven_flash_v2", "rime/arcana"). */
  model: string;
  /** Voice to use for the model. */
  voice: string;
  /** Extra configuration for the model. */
  extraKwargs?: Record<string, unknown>;
}

export type TTSFallbackModelType = TTSFallbackModel | string;

/** Normalize a single or list of FallbackModelType into TTSFallbackModel[]. */
export function normalizeTTSFallback(
  fallback: TTSFallbackModelType | TTSFallbackModelType[],
): TTSFallbackModel[] {
  const makeFallback = (model: TTSFallbackModelType): TTSFallbackModel => {
    if (typeof model === 'string') {
      const [name, voice] = parseTTSModelString(model);
      return { model: name, voice: voice ?? '' };
    }
    return model;
  };

  if (Array.isArray(fallback)) {
    return fallback.map(makeFallback);
  }
  return [makeFallback(fallback)];
}

type TTSEncoding = 'pcm_s16le';

const DEFAULT_ENCODING: TTSEncoding = 'pcm_s16le';
const DEFAULT_SAMPLE_RATE = 16000;
const NUM_CHANNELS = 1;
const DEFAULT_LANGUAGE = 'en';

export interface InferenceTTSOptions<TModel extends TTSModels> {
  model?: TModel;
  voice?: string;
  language?: LanguageCode;
  encoding: TTSEncoding;
  sampleRate: number;
  baseURL: string;
  apiKey: string;
  apiSecret: string;
  /** Flat provider-specific inference options forwarded as the `extra` payload field. */
  modelOptions: TTSOptions<TModel>;
  fallback?: TTSFallbackModel[];
  connOptions?: APIConnectOptions;
}

/**
 * Livekit Cloud Inference TTS
 */
export class TTS<TModel extends TTSModels> extends BaseTTS {
  private opts: InferenceTTSOptions<TModel>;
  private streams: Set<SynthesizeStream<TModel>> = new Set();
  pool: ConnectionPool<WebSocket>;

  #logger = log();
  #alignedTranscript: boolean;

  constructor(opts: {
    model: TModel;
    voice?: string;
    language?: string;
    baseURL?: string;
    encoding?: TTSEncoding;
    sampleRate?: number;
    apiKey?: string;
    apiSecret?: string;
    /** Flat provider-specific inference options forwarded as the `extra` payload field. */
    modelOptions?: TTSOptions<TModel>;
    fallback?: TTSFallbackModelType | TTSFallbackModelType[];
    connOptions?: APIConnectOptions;
  }) {
    const sampleRate = opts?.sampleRate ?? DEFAULT_SAMPLE_RATE;
    const initialAligned = hasAlignedTranscript(
      opts?.model,
      opts?.modelOptions as Record<string, unknown> | undefined,
    );
    super(sampleRate, 1, { streaming: true, alignedTranscript: initialAligned });

    const {
      model,
      voice,
      language = DEFAULT_LANGUAGE,
      baseURL,
      encoding = DEFAULT_ENCODING,
      apiKey,
      apiSecret,
      modelOptions = {} as TTSOptions<TModel>,
      fallback,
      connOptions,
    } = opts || {};

    const lkBaseURL = baseURL || getDefaultInferenceUrl();
    const lkApiKey = apiKey || process.env.LIVEKIT_INFERENCE_API_KEY || process.env.LIVEKIT_API_KEY;
    if (!lkApiKey) {
      throw new Error('apiKey is required: pass apiKey or set LIVEKIT_API_KEY');
    }

    const lkApiSecret =
      apiSecret || process.env.LIVEKIT_INFERENCE_API_SECRET || process.env.LIVEKIT_API_SECRET;
    if (!lkApiSecret) {
      throw new Error('apiSecret is required: pass apiSecret or set LIVEKIT_API_SECRET');
    }

    // read voice id from the model if provided: "provider/model:voice_id"
    let nextModel = model;
    let nextVoice = voice;
    if (typeof nextModel === 'string') {
      const idx = nextModel.lastIndexOf(':');
      if (idx !== -1) {
        const voiceFromModel = nextModel.slice(idx + 1);
        if (nextVoice && nextVoice !== voiceFromModel) {
          this.#logger.warn(
            '`voice` is provided via both argument and model, using the one from the argument',
            { voice: nextVoice, model: nextModel },
          );
        } else {
          nextVoice = voiceFromModel;
        }
        nextModel = nextModel.slice(0, idx) as TModel;
      }
    }

    const normalizedFallback = fallback ? normalizeTTSFallback(fallback) : undefined;

    this.opts = {
      model: nextModel,
      voice: nextVoice,
      language: normalizeLanguage(language),
      encoding,
      sampleRate,
      baseURL: lkBaseURL,
      apiKey: lkApiKey,
      apiSecret: lkApiSecret,
      modelOptions,
      fallback: normalizedFallback,
      connOptions: connOptions ?? DEFAULT_API_CONNECT_OPTIONS,
    };

    this.#alignedTranscript = hasAlignedTranscript(
      nextModel,
      modelOptions as Record<string, unknown>,
    );

    // Initialize connection pool
    this.pool = new ConnectionPool<WebSocket>({
      connectCb: (timeout) => this.connectWs(timeout),
      closeCb: (ws) => this.closeWs(ws),
      maxSessionDuration: 300_000,
      markRefreshedOnGet: true,
      connectTimeout: 10_000, // 10 seconds default
    });
  }

  get label() {
    return 'inference.TTS';
  }

  get model(): string {
    return this.opts.model ?? 'unknown';
  }

  get provider(): string {
    return 'livekit';
  }

  override get capabilities() {
    return { streaming: true, alignedTranscript: this.#alignedTranscript };
  }

  static fromModelString(modelString: string): TTS<AnyString> {
    const [model, voice] = parseTTSModelString(modelString);
    return new TTS({ model, voice: voice || undefined });
  }

  updateOptions(
    opts: Partial<
      Pick<InferenceTTSOptions<TModel>, 'model' | 'voice' | 'language' | 'modelOptions'>
    >,
  ) {
    const mergedModelOptions =
      opts.modelOptions !== undefined
        ? ({ ...this.opts.modelOptions, ...opts.modelOptions } as TTSOptions<TModel>)
        : this.opts.modelOptions;

    const sessionAffectingChange =
      opts.model !== undefined ||
      opts.voice !== undefined ||
      opts.language !== undefined ||
      opts.modelOptions !== undefined;

    this.opts = {
      ...this.opts,
      ...opts,
      modelOptions: mergedModelOptions,
      language: opts.language !== undefined ? normalizeLanguage(opts.language) : this.opts.language,
    };

    this.#alignedTranscript = hasAlignedTranscript(
      this.opts.model,
      this.opts.modelOptions as Record<string, unknown>,
    );

    // The inference gateway only consumes `model` / `voice` / `language` /
    // `extra` (modelOptions) when the WebSocket is first opened via
    // `session.create`. Pooled sockets keep the old session settings and, for
    // example, would keep reporting `alignedTranscript=true` while the server
    // never emits `output_alignment`. Invalidate so the next `stream()` opens
    // a fresh connection with the up-to-date payload.
    if (sessionAffectingChange) {
      this.pool.invalidate();
    }

    for (const stream of this.streams) {
      stream.updateOptions(this.opts);
    }
  }

  synthesize(_: string): ChunkedStream {
    throw new Error('ChunkedStream is not implemented');
  }

  stream(options?: { connOptions?: APIConnectOptions }): SynthesizeStream<TModel> {
    const { connOptions = this.opts.connOptions ?? DEFAULT_API_CONNECT_OPTIONS } = options || {};
    const stream = new SynthesizeStream(this, { ...this.opts }, connOptions);
    this.streams.add(stream);
    return stream;
  }

  async connectWs(timeout: number): Promise<WebSocket> {
    let baseURL = this.opts.baseURL;
    if (baseURL.startsWith('http://') || baseURL.startsWith('https://')) {
      baseURL = baseURL.replace('http', 'ws');
    }

    const token = await createAccessToken(this.opts.apiKey, this.opts.apiSecret);
    const url = `${baseURL}/tts`;
    const headers = { Authorization: `Bearer ${token}` } as Record<string, string>;

    const params = {
      type: 'session.create',
      sample_rate: String(this.opts.sampleRate),
      encoding: this.opts.encoding,
      extra: this.opts.modelOptions,
    } as Record<string, unknown>;

    if (this.opts.voice) (params as Record<string, unknown>).voice = this.opts.voice;
    if (this.opts.model) (params as Record<string, unknown>).model = this.opts.model;
    if (this.opts.language) (params as Record<string, unknown>).language = this.opts.language;

    if (this.opts.fallback?.length) {
      params.fallback = {
        models: this.opts.fallback.map((m) => ({
          model: m.model,
          voice: m.voice,
          extra: m.extraKwargs ?? {},
        })),
      };
    }

    if (this.opts.connOptions) {
      params.connection = {
        timeout: this.opts.connOptions.timeoutMs / 1000,
        retries: this.opts.connOptions.maxRetry,
      };
    }

    this.#logger.debug({ url }, 'inference.TTS creating new websocket connection (pool miss)');
    const socket = await connectWs(url, headers, timeout);
    socket.send(JSON.stringify(params));
    return socket;
  }

  async closeWs(ws: WebSocket) {
    await ws.close();
  }

  prewarm(): void {
    this.pool.prewarm();
  }

  async close() {
    for (const stream of this.streams) {
      await stream.close();
    }
    this.streams.clear();
    await this.pool.close();
  }
}

export class SynthesizeStream<TModel extends TTSModels> extends BaseSynthesizeStream {
  private opts: InferenceTTSOptions<TModel>;
  private tts: TTS<TModel>;

  #logger = log();

  constructor(tts: TTS<TModel>, opts: InferenceTTSOptions<TModel>, connOptions: APIConnectOptions) {
    super(tts, connOptions);
    this.opts = opts;
    this.tts = tts;
  }

  get label() {
    return 'inference.SynthesizeStream';
  }

  updateOptions(
    opts: Partial<
      Pick<InferenceTTSOptions<TModel>, 'model' | 'voice' | 'language' | 'modelOptions'>
    >,
  ) {
    const mergedModelOptions =
      opts.modelOptions !== undefined
        ? ({ ...this.opts.modelOptions, ...opts.modelOptions } as TTSOptions<TModel>)
        : this.opts.modelOptions;

    this.opts = {
      ...this.opts,
      ...opts,
      modelOptions: mergedModelOptions,
      language: opts.language !== undefined ? normalizeLanguage(opts.language) : this.opts.language,
    };
  }

  protected async run(): Promise<void> {
    let closing = false;
    let lastFrame: AudioFrame | undefined;
    // Timestamps are delivered in their own WS message; buffer them and attach
    // to the next audio frame that we forward to the output emitter. This
    // mirrors the semantics of `output_emitter.push_timed_transcript` on the
    // Python side.
    let pendingTimedTranscripts: TimedString[] = [];

    const sendTokenizerStream = new tokenizeBasic.SentenceTokenizer().stream();
    const eventChannel = createStreamChannel<TtsServerEvent>();
    const requestId = shortuuid('tts_request_');
    const inputSentEvent = new Event();

    // Signal for protocol-driven completion (when 'done' message is received)
    const completionFuture = new Future<void>();

    const resourceCleanup = async () => {
      if (closing) return;
      closing = true;
      sendTokenizerStream.close();
      // close() returns a promise; don't leak it
      await eventChannel.close();
    };

    const sendClientEvent = async (event: TtsClientEvent, ws: WebSocket, signal: AbortSignal) => {
      // Don't send events to a closed WebSocket or aborted controller
      if (signal.aborted || closing) return;

      const validatedEvent = await ttsClientEventSchema.parseAsync(event);
      if (ws.readyState !== WebSocket.OPEN) {
        this.#logger.warn('Trying to send client TTS event to a closed WebSocket');
        return;
      }
      ws.send(JSON.stringify(validatedEvent));
    };

    const sendLastFrame = (segmentId: string, final: boolean) => {
      if (lastFrame) {
        this.queue.put({
          requestId,
          segmentId,
          frame: lastFrame,
          final,
          timedTranscripts:
            pendingTimedTranscripts.length > 0 ? pendingTimedTranscripts : undefined,
        });
        lastFrame = undefined;
        pendingTimedTranscripts = [];
      }
    };

    const createInputTask = async (signal: AbortSignal) => {
      for await (const data of this.input) {
        if (signal.aborted || closing) break;
        if (data === SynthesizeStream.FLUSH_SENTINEL) {
          sendTokenizerStream.flush();
          continue;
        }
        sendTokenizerStream.pushText(data);
      }
      // Only call endInput if the stream hasn't been closed by cleanup
      if (!closing) {
        sendTokenizerStream.endInput();
      }
    };

    const createSentenceStreamTask = async (ws: WebSocket, signal: AbortSignal) => {
      for await (const ev of sendTokenizerStream) {
        if (signal.aborted || closing) break;

        // Carry per-utterance generation config so mid-stream voice/model/language
        // changes ride the gateway's hot path (no reconnect needed on the active
        // session). The gateway also merges `extra` into the live session.
        const generationConfig: Record<string, string> = {};
        if (this.opts.voice) generationConfig.voice = this.opts.voice;
        if (this.opts.model) generationConfig.model = this.opts.model;
        if (this.opts.language) generationConfig.language = this.opts.language;

        this.markStarted();
        await sendClientEvent(
          {
            type: 'input_transcript',
            transcript: ev.token + ' ',
            generation_config: generationConfig,
            extra: (this.opts.modelOptions as Record<string, unknown>) ?? {},
          },
          ws,
          signal,
        );
        inputSentEvent.set();
      }

      await sendClientEvent({ type: 'session.flush' }, ws, signal);
      // needed in case empty input is sent
      inputSentEvent.set();
    };

    // Handles WebSocket message routing and error handling
    // Completes based on protocol messages, NOT on ws.close()
    const createWsListenerTask = async (ws: WebSocket, signal: AbortSignal) => {
      const onMessage = (data: Buffer) => {
        try {
          const eventJson = JSON.parse(data.toString()) as Record<string, unknown>;
          const validatedEvent = ttsServerEventSchema.parse(eventJson);
          // writer.write returns a promise; avoid unhandled rejections if stream is closed
          void eventChannel.write(validatedEvent).catch((error) => {
            this.#logger.debug(
              { error },
              'Failed writing TTS event to stream channel (likely closed)',
            );
          });
        } catch (e) {
          this.#logger.error({ error: e }, 'Error parsing WebSocket message');
        }
      };

      const onError = (e: Error) => {
        this.#logger.error({ error: e }, 'WebSocket error');
        void resourceCleanup();
        try {
          // If the ws is misbehaving, hard-stop it immediately to avoid buffering.
          ws.terminate?.();
        } catch {
          // ignore
        }
        // Ensure this ws is not reused
        this.tts.pool.remove(ws);
        completionFuture.reject(e);
      };

      const onClose = () => {
        // WebSocket closed unexpectedly (not by us)
        if (!closing) {
          this.#logger.error('WebSocket closed unexpectedly');
          void resourceCleanup();
          // Ensure this ws is not reused
          this.tts.pool.remove(ws);
          completionFuture.reject(
            new APIStatusError({
              message: 'Gateway connection closed unexpectedly',
              options: { requestId },
            }),
          );
        }
      };

      const onAbort = () => {
        void resourceCleanup();
        try {
          // On interruption/abort, close the websocket immediately so the server stops streaming
          // and the ws library doesn't buffer unread frames in memory.
          ws.terminate?.();
        } catch {
          // ignore
        }
        this.tts.pool.remove(ws);
        inputSentEvent.set();
        completionFuture.resolve();
      };

      // Attach listeners
      ws.on('message', onMessage);
      ws.on('error', onError);
      ws.on('close', onClose);
      signal.addEventListener('abort', onAbort);

      try {
        // Wait for protocol-driven completion or error
        await completionFuture.await;
      } finally {
        // IMPORTANT: Remove listeners so connection can be reused
        ws.off('message', onMessage);
        ws.off('error', onError);
        ws.off('close', onClose);
        signal.removeEventListener('abort', onAbort);
      }
    };

    const createRecvTask = async (signal: AbortSignal) => {
      let currentSessionId: string | null = null;
      const recvTimeoutMs = this.connOptions.timeoutMs;

      const bstream = new AudioByteStream(this.opts.sampleRate, NUM_CHANNELS);
      const serverEventStream = eventChannel.stream();
      const reader = serverEventStream.getReader();

      try {
        await inputSentEvent.wait();

        while (!this.closed && !signal.aborted) {
          const result = await waitUntilTimeout(
            reader.read(),
            recvTimeoutMs,
            () => new APITimeoutError({ message: 'TTS recv idle timeout' }),
          );

          if (signal.aborted) return;
          if (result.done) return;

          const parsedServerEvent = ttsKnownServerEventSchema.safeParse(result.value);
          if (!parsedServerEvent.success) {
            continue;
          }

          const serverEvent = parsedServerEvent.data;
          const sessionIdFromEvent = (serverEvent as { session_id?: string }).session_id;
          if (currentSessionId === null && sessionIdFromEvent) {
            currentSessionId = sessionIdFromEvent;
          }

          switch (serverEvent.type) {
            case 'session.created':
              currentSessionId = serverEvent.session_id;
              break;
            case 'output_audio':
              const base64Data = new Int8Array(Buffer.from(serverEvent.audio, 'base64'));
              for (const frame of bstream.write(base64Data.buffer)) {
                sendLastFrame(currentSessionId!, false);
                lastFrame = frame;
              }
              break;
            case 'output_alignment':
              if (serverEvent.words && serverEvent.words.length > 0) {
                for (const w of serverEvent.words) {
                  pendingTimedTranscripts.push(
                    createTimedString({
                      text: w.word,
                      startTime: w.start,
                      endTime: w.end,
                    }),
                  );
                }
              } else if (serverEvent.chars && serverEvent.chars.length > 0) {
                for (const c of serverEvent.chars) {
                  pendingTimedTranscripts.push(
                    createTimedString({
                      text: c.char,
                      startTime: c.start,
                      endTime: c.end,
                    }),
                  );
                }
              }
              break;
            case 'done':
              for (const frame of bstream.flush()) {
                sendLastFrame(currentSessionId!, false);
                lastFrame = frame;
              }
              sendLastFrame(currentSessionId!, true);
              this.queue.put(SynthesizeStream.END_OF_STREAM);
              await resourceCleanup();
              completionFuture.resolve();
              return;
            case 'session.closed':
              await resourceCleanup();
              completionFuture.resolve();
              return;
            case 'error':
              this.#logger.error(
                { serverEvent },
                'Received error message from LiveKit TTS WebSocket',
              );
              await resourceCleanup();
              completionFuture.reject(
                new APIError(`LiveKit TTS returned error: ${serverEvent.message}`),
              );
              return;
          }
        }
      } catch (e) {
        if (e instanceof APITimeoutError) {
          this.#logger.warn('TTS recv task timed out waiting for server message');
          await resourceCleanup();
          completionFuture.reject(e);
          return;
        }
        throw e;
      } finally {
        reader.releaseLock();
        try {
          await serverEventStream.cancel();
        } catch (e) {
          this.#logger.debug('Error cancelling serverEventStream (may already be cancelled):', e);
        }
      }
    };

    try {
      await this.tts.pool.withConnection(
        async (ws: WebSocket) => {
          try {
            // IMPORTANT: don't cancel the stream's controller on normal completion,
            // otherwise the pool will remove+close the ws and every run becomes a pool miss.
            const runController = new AbortController();
            const onStreamAbort = () => runController.abort(this.abortController.signal.reason);
            this.abortController.signal.addEventListener('abort', onStreamAbort, { once: true });

            const tasks = [
              Task.from(
                async (controller) => {
                  const combined = combineSignals(runController.signal, controller.signal);
                  await createInputTask(combined);
                },
                undefined,
                'inference-tts-input',
              ),
              Task.from(
                async (controller) => {
                  const combined = combineSignals(runController.signal, controller.signal);
                  await createSentenceStreamTask(ws, combined);
                },
                undefined,
                'inference-tts-sentence',
              ),
              Task.from(
                async (controller) => {
                  const combined = combineSignals(runController.signal, controller.signal);
                  await createWsListenerTask(ws, combined);
                },
                undefined,
                'inference-tts-ws-listener',
              ),
              Task.from(
                async (controller) => {
                  const combined = combineSignals(runController.signal, controller.signal);
                  await createRecvTask(combined);
                },
                undefined,
                'inference-tts-recv',
              ),
            ];

            try {
              await ThrowsPromise.all(tasks.map((t) => t.result));
            } finally {
              // Mirror python finally: unblock recv and cancel all tasks.
              inputSentEvent.set();
              await resourceCleanup();
              await cancelAndWait(tasks, 5000);
              this.abortController.signal.removeEventListener('abort', onStreamAbort);
            }
          } catch (e) {
            // If aborted, don't throw - let cleanup handle it
            if (e instanceof Error && e.name === 'AbortError') {
              return;
            }
            throw e;
          }
        },
        {
          timeout: this.connOptions.timeoutMs,
        },
      );
    } catch (e) {
      // Handle connection errors
      if (e instanceof Error && e.name === 'AbortError') {
        // Abort is expected during normal shutdown
        return;
      }
      throw e;
    } finally {
      // Ensure cleanup always runs (and don't leak the promise)
      await resourceCleanup();
    }
  }
}
