// SPDX-FileCopyrightText: 2026 LiveKit, Inc.
//
// SPDX-License-Identifier: Apache-2.0
//
import {
  type APIConnectOptions,
  type AudioBuffer,
  AudioByteStream,
  ChatMessage,
  type ConversationItemAddedEvent,
  Future,
  Task,
  createTimedString,
  delay,
  log,
  normalizeLanguage,
  stt,
  waitForAbort,
} from '@livekit/agents';
import type { AudioFrame } from '@livekit/rtc-node';
import type { RawData } from 'ws';
import { WebSocket } from 'ws';
import type { STTEncoding, STTModels, VoiceFocus } from './models.js';

// Speech models in the Universal-3 Pro family, which share the same parameter support.
const U3_PRO_MODELS = ['u3-rt-pro', 'u3-rt-pro-beta-1', 'universal-3-5-pro'] as const;

function isU3ProModel(model: STTModels): boolean {
  return U3_PRO_MODELS.includes(model as (typeof U3_PRO_MODELS)[number]);
}

// AssemblyAI Universal-Streaming (v3) message envelope. All fields are optional
// since we narrow on `type` before reading anything else.
interface StreamEventMessage {
  type?: 'Begin' | 'SpeechStarted' | 'Turn' | 'Termination' | string;
  // Begin
  id?: string;
  expires_at?: number;
  // Turn
  transcript?: string;
  utterance?: string;
  end_of_turn?: boolean;
  end_of_turn_confidence?: number;
  turn_is_formatted?: boolean;
  language_code?: string;
  language_confidence?: number;
  speaker_label?: string;
  words?: Array<{
    text?: string;
    start?: number;
    end?: number;
    confidence?: number;
    speaker?: string;
  }>;
  // Termination
  audio_duration_seconds?: number;
  session_duration_seconds?: number;
}

function speechDataMetadata(data: StreamEventMessage): stt.SpeechData['metadata'] | undefined {
  const assemblyai: Record<string, number> = {};

  if (typeof data.language_confidence === 'number') {
    assemblyai.languageConfidence = data.language_confidence;
  }

  if (Object.keys(assemblyai).length === 0) return undefined;

  return { assemblyai };
}

export interface STTOptions {
  apiKey?: string;
  sampleRate: number;
  /**
   * How large each chunk of audio is before being sent to AssemblyAI, in
   * milliseconds. Corresponds to Python's `buffer_size_seconds` (seconds there,
   * ms here per this repo's time-unit convention).
   */
  bufferSizeMs: number;
  encoding: STTEncoding;
  speechModel: STTModels;
  languageDetection?: boolean;
  /**
   * Session inactivity timeout in seconds. AssemblyAI accepts integer values
   * from 5 to 3600; when unset, no inactivity timeout is applied.
   */
  inactivityTimeout?: number;
  endOfTurnConfidenceThreshold?: number;
  /** Minimum silence (ms) before a confident end-of-turn is finalized. */
  minTurnSilence?: number;
  /** Maximum silence (ms) before end-of-turn is forced regardless of confidence. */
  maxTurnSilence?: number;
  formatTurns?: boolean;
  keytermsPrompt?: string[];
  /** Only supported with the Universal-3 Pro model family. */
  prompt?: string;
  /** Only supported with the Universal-3 Pro model family. */
  agentContext?: string;
  /** Only supported with the Universal-3 Pro model family. Set at connection time only. */
  previousContextNTurns?: number;
  vadThreshold?: number;
  /**
   * Enable speaker diarization. Note: AssemblyAI will return per-word speaker
   * labels, but the JS framework's `stt.SpeechData` type does not yet expose
   * a `speakerId` field (unlike the Python framework), so the labels are not
   * currently surfaced on emitted events. Setting this to `true` still has
   * effect server-side. Once the base `SpeechData` interface gains speaker
   * support, `#processStreamEvent` should forward `data.words[].speaker` too.
   */
  speakerLabels?: boolean;
  maxSpeakers?: number;
  domain?: string;
  /** Isolate the primary voice and suppress background noise. Connect-time only. */
  voiceFocus?: VoiceFocus;
  /** Background audio suppression aggressiveness, from 0.0 to 1.0. Connect-time only. */
  voiceFocusThreshold?: number;
  /**
   * Accuracy/latency preset for the Universal-3 Pro model family: `min_latency`, `balanced`,
   * or `max_accuracy`. Explicit turn-silence values still take precedence over mode defaults.
   */
  mode?: 'min_latency' | 'balanced' | 'max_accuracy';
  /**
   * When the model supports it, let an `AgentSession` push each assistant reply into
   * `agentContext` so it is carried into the model's conversation context. Defaults to false;
   * set true to enable. Prior user turns are carried automatically by the model regardless of
   * this flag. Ignored on models without context support.
   */
  agentContextCarryover?: boolean;
  baseUrl: string;
}

const defaultSTTOptions: STTOptions = {
  apiKey: process.env.ASSEMBLYAI_API_KEY,
  sampleRate: 16000,
  bufferSizeMs: 50,
  encoding: 'pcm_s16le',
  speechModel: 'universal-3-5-pro',
  baseUrl: 'wss://streaming.assemblyai.com',
};

export class STT extends stt.STT {
  #opts: STTOptions;
  #streams = new Set<WeakRef<SpeechStream>>();
  // set (user + session))
  #userKeyterms: string[];
  #sessionKeyterms: string[] = [];
  label = 'assemblyai.STT';

  get model(): string {
    return this.#opts.speechModel;
  }

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

  constructor(opts: Partial<STTOptions> = {}) {
    // u3-rt-pro family — "u3-pro" is normalized below — and is opt-in via the user)
    const rawModel = opts.speechModel ?? defaultSTTOptions.speechModel;
    const supportsCarryover = isU3ProModel(rawModel) || rawModel === 'u3-pro';
    if (opts.agentContextCarryover && !supportsCarryover) {
      log().warn(
        `agentContextCarryover is enabled but model '${rawModel}' does not support it; ignoring`,
      );
    }
    super({
      streaming: true,
      interimResults: true,
      alignedTranscript: 'word',
      keyterms: true,
      chatContext: (opts.agentContextCarryover ?? false) && supportsCarryover,
    });

    if (opts.speechModel === 'u3-pro') {
      log().warn("'u3-pro' is deprecated, use 'universal-3-5-pro' instead.");
      opts.speechModel = 'universal-3-5-pro';
    }

    const speechModel = opts.speechModel ?? defaultSTTOptions.speechModel;
    if (!isU3ProModel(speechModel)) {
      for (const param of [
        'prompt',
        'agentContext',
        'previousContextNTurns',
        'voiceFocus',
        'voiceFocusThreshold',
        'mode',
      ] as const) {
        if (opts[param] !== undefined) {
          throw new Error(
            `The '${param}' parameter is only supported with the ${U3_PRO_MODELS.join(', ')} models.`,
          );
        }
      }
    }

    const apiKey = opts.apiKey ?? defaultSTTOptions.apiKey;
    if (!apiKey) {
      throw new Error(
        'AssemblyAI API key is required. Pass one in via the `apiKey` parameter, or set it as the `ASSEMBLYAI_API_KEY` environment variable',
      );
    }

    // Minimize latency by default, but let AssemblyAI's mode preset control silence tuning.
    const minTurnSilence = opts.minTurnSilence ?? (opts.mode === undefined ? 100 : undefined);

    this.#opts = {
      ...defaultSTTOptions,
      ...opts,
      apiKey,
      minTurnSilence,
    };
    this.#userKeyterms = [...(this.#opts.keytermsPrompt ?? [])];
  }

  // eslint-disable-next-line @typescript-eslint/no-unused-vars
  async _recognize(_: AudioBuffer): Promise<stt.SpeechEvent> {
    throw new Error('Non-streaming recognize is not supported on AssemblyAI STT');
  }

  updateOptions(opts: Partial<STTOptions>) {
    // session keyterms so a user update doesn't drop them)
    const nextOpts = { ...opts };
    if (nextOpts.keytermsPrompt !== undefined) {
      this.#userKeyterms = [...nextOpts.keytermsPrompt];
      nextOpts.keytermsPrompt = [...new Set([...this.#userKeyterms, ...this.#sessionKeyterms])];
    }
    this.#opts = { ...this.#opts, ...nextOpts };
    for (const ref of this.#streams) {
      const stream = ref.deref();
      if (stream) {
        stream.updateOptions(nextOpts);
      } else {
        this.#streams.delete(ref);
      }
    }
  }

  override _updateSessionKeyterms(keyterms: string[]): void {
    if (
      keyterms.length === this.#sessionKeyterms.length &&
      keyterms.every((t, i) => t === this.#sessionKeyterms[i])
    ) {
      return;
    }
    this.#sessionKeyterms = [...keyterms];
    const merged = [...new Set([...this.#userKeyterms, ...keyterms])];
    this.#opts.keytermsPrompt = merged;
    // applied live via the stream's UpdateConfiguration (no reconnect)
    for (const ref of this.#streams) {
      const stream = ref.deref();
      if (stream) {
        stream.updateOptions({ keytermsPrompt: merged });
      } else {
        this.#streams.delete(ref);
      }
    }
  }

  override _pushConversationItem(ev: ConversationItemAddedEvent): void {
    const chatItem = ev.item;
    if (chatItem instanceof ChatMessage && chatItem.role === 'assistant' && chatItem.textContent) {
      this.updateOptions({ agentContext: chatItem.textContent });
    }
  }

  stream(options?: { connOptions?: APIConnectOptions }): SpeechStream {
    const stream = new SpeechStream(this, this.#opts, options?.connOptions);
    this.#streams.add(new WeakRef(stream));
    return stream;
  }
}

export class SpeechStream extends stt.SpeechStream {
  static readonly CLOSE_MSG = JSON.stringify({ type: 'Terminate' });

  #opts: STTOptions;
  #logger = log();
  #speechDurationInS = 0;
  #lastPreflightStartTime = 0;
  #pendingConfigMessages: Record<string, unknown>[] = [];
  #configMessagePending = new Future();
  #sessionId: string | null = null;
  #expiresAt: number | null = null;
  label = 'assemblyai.SpeechStream';

  constructor(stt: STT, opts: STTOptions, connOptions?: APIConnectOptions) {
    super(stt, opts.sampleRate, connOptions);
    this.#opts = opts;
    this.closed = false;
  }

  /**
   * The AssemblyAI session ID. Set when the WebSocket connection is established
   * (before any speech events). Null until the connection completes.
   * Share this with the AssemblyAI team when reporting issues.
   */
  get sessionId(): string | null {
    return this.#sessionId;
  }

  /**
   * Unix timestamp when the AssemblyAI session expires. Set alongside
   * {@link sessionId} when the WebSocket connection is established.
   */
  get expiresAt(): number | null {
    return this.#expiresAt;
  }

  updateOptions(opts: Partial<STTOptions>) {
    this.#opts = { ...this.#opts, ...opts };

    const configMsg: Record<string, unknown> = { type: 'UpdateConfiguration' };
    if (opts.prompt !== undefined) configMsg.prompt = opts.prompt;
    if (opts.agentContext !== undefined) configMsg.agent_context = opts.agentContext;
    if (opts.keytermsPrompt !== undefined) configMsg.keyterms_prompt = opts.keytermsPrompt;
    if (opts.maxTurnSilence !== undefined) configMsg.max_turn_silence = opts.maxTurnSilence;
    if (opts.minTurnSilence !== undefined) configMsg.min_turn_silence = opts.minTurnSilence;
    if (opts.endOfTurnConfidenceThreshold !== undefined) {
      configMsg.end_of_turn_confidence_threshold = opts.endOfTurnConfidenceThreshold;
    }
    if (opts.vadThreshold !== undefined) configMsg.vad_threshold = opts.vadThreshold;

    // Only send if any actual fields (besides `type`) were specified.
    if (Object.keys(configMsg).length > 1) {
      this.#pendingConfigMessages.push(configMsg);
      if (!this.#configMessagePending.done) this.#configMessagePending.resolve();
    }
  }

  /**
   * Force-finalize the current turn immediately.
   */
  forceEndpoint() {
    this.#pendingConfigMessages.push({ type: 'ForceEndpoint' });
    if (!this.#configMessagePending.done) this.#configMessagePending.resolve();
  }

  // Deepgram-style reconnect loop around a single websocket lifetime.
  protected async run() {
    const maxRetry = 32;
    let retries = 0;

    while (!this.input.closed && !this.closed) {
      try {
        const ws = await this.#connectWS();
        await this.#runWS(ws);
        retries = 0;
      } catch (e) {
        if (!this.closed && !this.input.closed) {
          if (retries >= maxRetry) {
            throw new Error(`failed to connect to AssemblyAI after ${retries} attempts: ${e}`);
          }

          const retryDelaySeconds = Math.min(retries * 5, 10);
          retries++;

          this.#logger.warn(
            `failed to connect to AssemblyAI, retrying in ${retryDelaySeconds} seconds: ${e} (${retries}/${maxRetry})`,
          );
          await delay(retryDelaySeconds * 1000);
        } else {
          this.#logger.warn(
            `AssemblyAI disconnected, connection is closed: ${e} (inputClosed: ${this.input.closed}, isClosed: ${this.closed})`,
          );
        }
      }
    }

    this.closed = true;
  }

  async #connectWS(): Promise<WebSocket> {
    // Universal-3 Pro family models default both min and max silence to 100ms when unset.
    // When a mode preset is selected, leave them unset unless explicitly provided so the
    // server's per-mode silence tuning is not overridden by the latency-optimized default.
    let minSilence = this.#opts.minTurnSilence;
    let maxSilence = this.#opts.maxTurnSilence;
    if (isU3ProModel(this.#opts.speechModel)) {
      if (minSilence === undefined && this.#opts.mode === undefined) minSilence = 100;
      if (maxSilence === undefined) maxSilence = minSilence;
    }

    // Default language_detection to true for multilingual / u3-rt-pro-family models, false otherwise.
    const defaultLanguageDetection =
      this.#opts.speechModel.includes('multilingual') || isU3ProModel(this.#opts.speechModel);
    const languageDetection = this.#opts.languageDetection ?? defaultLanguageDetection;

    const liveConfig: Record<string, unknown> = {
      sample_rate: this.#opts.sampleRate,
      encoding: this.#opts.encoding,
      speech_model: this.#opts.speechModel,
      format_turns: this.#opts.formatTurns,
      end_of_turn_confidence_threshold: this.#opts.endOfTurnConfidenceThreshold,
      min_turn_silence: minSilence,
      max_turn_silence: maxSilence,
      keyterms_prompt:
        this.#opts.keytermsPrompt !== undefined && this.#opts.keytermsPrompt.length > 0
          ? JSON.stringify(this.#opts.keytermsPrompt)
          : undefined,
      language_detection: languageDetection,
      inactivity_timeout: this.#opts.inactivityTimeout,
      prompt: this.#opts.prompt,
      agent_context: this.#opts.agentContext,
      previous_context_n_turns: this.#opts.previousContextNTurns,
      vad_threshold: this.#opts.vadThreshold,
      speaker_labels: this.#opts.speakerLabels,
      max_speakers: this.#opts.maxSpeakers,
      domain: this.#opts.domain,
      voice_focus: this.#opts.voiceFocus,
      voice_focus_threshold: this.#opts.voiceFocusThreshold,
      mode: this.#opts.mode,
    };

    const url = new URL(`${this.#opts.baseUrl}/v3/ws`);
    // Python serializes booleans as the strings "true"/"false", so we mirror that.
    for (const [key, value] of Object.entries(liveConfig)) {
      if (value === undefined || value === null) continue;
      if (typeof value === 'boolean') {
        url.searchParams.append(key, value ? 'true' : 'false');
      } else {
        url.searchParams.append(key, String(value));
      }
    }

    const ws = new WebSocket(url, {
      headers: {
        Authorization: this.#opts.apiKey!,
        'Content-Type': 'application/json',
        'User-Agent': 'AssemblyAI/1.0 (integration=Livekit)',
      },
    });

    await new Promise<void>((resolve, reject) => {
      ws.on('open', () => resolve());
      ws.on('error', (error) => reject(error));
      ws.on('close', (code) => reject(new Error(`WebSocket returned ${code}`)));
    });

    return ws;
  }

  async #runWS(ws: WebSocket) {
    let closing = false;
    const sessionController = new AbortController();

    // gets cancelled also when sendTask is complete
    const wsMonitor = Task.from(async (controller) => {
      const closed = new Promise<void>((_, reject) => {
        ws.once('close', (code, reason) => {
          if (!closing) {
            this.#logger.error(`WebSocket closed with code ${code}: ${reason}`);
            reject(new Error('WebSocket closed'));
          }
        });
      });

      await Promise.race([closed, waitForAbort(controller.signal)]);
    });

    const sendTask = async () => {
      const samplesPerBuffer = Math.floor((this.#opts.sampleRate * this.#opts.bufferSizeMs) / 1000);
      const audioStream = new AudioByteStream(this.#opts.sampleRate, 1, samplesPerBuffer);

      const abortPromise = waitForAbort(this.abortSignal);
      const sessionAbort = waitForAbort(sessionController.signal);

      try {
        while (!this.closed) {
          const result = await Promise.race([this.input.next(), abortPromise, sessionAbort]);

          if (result === undefined) return; // aborted
          if (result.done) break;

          const data = result.value;

          let frames: AudioFrame[];
          if (data === SpeechStream.FLUSH_SENTINEL) {
            frames = audioStream.flush();
          } else if (data.sampleRate === this.#opts.sampleRate && data.channels === 1) {
            // AssemblyAI expects mono PCM. The base SpeechStream only resamples
            // sample rate, so reject any frame that is not already downmixed.
            frames = audioStream.write(data.data.buffer as ArrayBuffer);
          } else {
            throw new Error('sample rate or channel count of frame does not match');
          }

          for (const frame of frames) {
            this.#speechDurationInS += frame.samplesPerChannel / frame.sampleRate;
            ws.send(frame.data.buffer);
          }
        }
      } finally {
        closing = true;
        try {
          ws.send(SpeechStream.CLOSE_MSG);
        } catch {
          // ignore — socket may already be closing
        }
        wsMonitor.cancel();
      }
    };

    let messageHandler: ((msg: RawData, isBinary: boolean) => void) | null = null;
    const listenTask = Task.from(async (controller) => {
      const listenMessage = new Promise<void>((resolve, reject) => {
        messageHandler = (msg, isBinary) => {
          if (isBinary) {
            this.#logger.error('unexpected binary message from AssemblyAI');
            return;
          }
          try {
            const json = JSON.parse(msg.toString()) as StreamEventMessage;
            this.#processStreamEvent(json);
            if (this.closed || closing) {
              resolve();
            }
          } catch (err) {
            this.#logger.error(`AssemblyAI: error processing message: ${msg}`);
            reject(err);
          }
        };
        ws.on('message', messageHandler);
      });

      await Promise.race([listenMessage, waitForAbort(controller.signal)]);
    });

    const configTask = Task.from(async (controller) => {
      // Drain any messages queued while the socket was reconnecting.
      while (this.#pendingConfigMessages.length > 0) {
        const msg = this.#pendingConfigMessages.shift()!;
        ws.send(JSON.stringify(msg));
      }

      while (!controller.signal.aborted) {
        await Promise.race([this.#configMessagePending.await, waitForAbort(controller.signal)]);
        if (controller.signal.aborted) return;

        this.#configMessagePending = new Future();
        while (this.#pendingConfigMessages.length > 0) {
          const msg = this.#pendingConfigMessages.shift()!;
          ws.send(JSON.stringify(msg));
        }
      }
    });

    try {
      await Promise.all([sendTask(), listenTask.result, wsMonitor.result]);
    } finally {
      closing = true;
      sessionController.abort();
      listenTask.cancel();
      configTask.cancel();
      if (messageHandler) ws.off('message', messageHandler);
      try {
        ws.close();
      } catch {
        // ignore
      }
    }
  }

  #averageConfidence(words: Array<{ confidence?: number }>): number {
    if (words.length === 0) return 0;
    return words.reduce((sum, w) => sum + (w.confidence ?? 0), 0) / words.length;
  }

  #processStreamEvent(data: StreamEventMessage) {
    const messageType = data.type;

    if (messageType === 'Begin') {
      this.#sessionId = data.id ?? null;
      this.#expiresAt = data.expires_at ?? null;
      this.#logger.info(
        `AssemblyAI session started id=${this.#sessionId} expires_at=${this.#expiresAt}`,
      );
      return;
    }

    if (messageType === 'SpeechStarted') {
      this.queue.put({ type: stt.SpeechEventType.START_OF_SPEECH });
      return;
    }

    if (messageType === 'Termination') {
      this.#logger.debug(
        `AssemblyAI session terminated audio_duration=${data.audio_duration_seconds}s session_duration=${data.session_duration_seconds}s`,
      );
      return;
    }

    if (messageType !== 'Turn') {
      return;
    }

    const words = data.words ?? [];
    const endOfTurn = Boolean(data.end_of_turn);
    const turnIsFormatted = Boolean(data.turn_is_formatted);
    const utterance = data.utterance ?? '';
    const transcript = data.transcript ?? '';
    const language = normalizeLanguage(data.language_code ?? 'en');
    const metadata = speechDataMetadata(data);

    // Word timestamps are in milliseconds:
    // https://www.assemblyai.com/docs/api-reference/streaming-api/streaming-api#receive.receiveTurn.words
    const timedWords = words.map((word) =>
      createTimedString({
        text: word.text ?? '',
        startTime: (word.start ?? 0) / 1000 + this.startTimeOffset,
        endTime: (word.end ?? 0) / 1000 + this.startTimeOffset,
        confidence: word.confidence ?? 0,
        startTimeOffset: this.startTimeOffset,
      }),
    );

    let startTime = 0;
    let endTime = 0;
    let confidence = 0;

    // `words` are cumulative for the turn — emit as an interim transcript.
    if (timedWords.length > 0) {
      const interimText = timedWords.map((w) => w.text).join(' ');
      startTime = timedWords[0]!.startTime ?? 0;
      endTime = timedWords[timedWords.length - 1]!.endTime ?? 0;
      confidence = this.#averageConfidence(timedWords);

      this.queue.put({
        type: stt.SpeechEventType.INTERIM_TRANSCRIPT,
        alternatives: [
          {
            language,
            text: interimText,
            startTime,
            endTime,
            confidence,
            words: timedWords,
            ...(metadata ? { metadata } : {}),
          },
        ],
      });
    }

    // `utterance` is chunk-based (not cumulative) — emit as a preflight transcript
    // covering only the words since the last preflight.
    if (utterance) {
      if (this.#lastPreflightStartTime === 0) {
        this.#lastPreflightStartTime = startTime;
      }

      const utteranceWords = timedWords.filter(
        (w) => w.startTime !== undefined && w.startTime >= this.#lastPreflightStartTime,
      );
      const utteranceConfidence = this.#averageConfidence(utteranceWords);

      this.queue.put({
        type: stt.SpeechEventType.PREFLIGHT_TRANSCRIPT,
        alternatives: [
          {
            language,
            text: utterance,
            startTime: this.#lastPreflightStartTime,
            endTime,
            confidence: utteranceConfidence,
            words: utteranceWords,
            ...(metadata ? { metadata } : {}),
          },
        ],
      });
      this.#lastPreflightStartTime = endTime;
    }

    // End-of-turn: emit FINAL_TRANSCRIPT + END_OF_SPEECH.
    // If the user asked for formatted turns, wait for a formatted final.
    const waitingForFormatted = this.#opts.formatTurns === true && !turnIsFormatted;
    if (endOfTurn && !waitingForFormatted) {
      this.queue.put({
        type: stt.SpeechEventType.FINAL_TRANSCRIPT,
        alternatives: [
          {
            language,
            text: transcript,
            startTime,
            endTime,
            confidence,
            words: timedWords,
            ...(metadata ? { metadata } : {}),
          },
        ],
      });

      this.queue.put({ type: stt.SpeechEventType.END_OF_SPEECH });
      if (this.#speechDurationInS > 0) {
        this.queue.put({
          type: stt.SpeechEventType.RECOGNITION_USAGE,
          // Propagate the AssemblyAI session id as the request id so metrics
          // can be correlated back to a specific connection, mirroring how
          // Deepgram surfaces its `request_id`.
          requestId: this.#sessionId ?? undefined,
          recognitionUsage: {
            audioDuration: this.#speechDurationInS,
          },
        });
        this.#speechDurationInS = 0;
        this.#lastPreflightStartTime = 0;
      }
    }
  }
}
