import { ThrowsPromise } from "@livekit/throws-transformer/throws";
import { TransformStream } from "stream/web";
import WebSocket from "ws";
import { z } from "zod";
import { APIConnectionError, APIStatusError, APITimeoutError } from "../../_exceptions.js";
import { log } from "../../log.js";
import { Event } from "../../utils.js";
import { buildMetadataHeaders, createAccessToken } from "../utils.js";
import { InterruptionCacheEntry } from "./interruption_cache_entry.js";
const MSG_SESSION_CREATE = "session.create";
const MSG_SESSION_CLOSE = "session.close";
const MSG_SESSION_CREATED = "session.created";
const MSG_SESSION_CLOSED = "session.closed";
const MSG_INTERRUPTION_DETECTED = "bargein_detected";
const MSG_INFERENCE_DONE = "inference_done";
const MSG_ERROR = "error";
const wsMessageSchema = z.discriminatedUnion("type", [
  z.object({
    type: z.literal(MSG_SESSION_CREATED),
    default_threshold: z.number().nullish()
  }),
  z.object({
    type: z.literal(MSG_SESSION_CLOSED)
  }),
  z.object({
    type: z.literal(MSG_INTERRUPTION_DETECTED),
    created_at: z.number(),
    probabilities: z.array(z.number()).default([]),
    prediction_duration: z.number().default(0)
  }),
  z.object({
    type: z.literal(MSG_INFERENCE_DONE),
    created_at: z.number(),
    probabilities: z.array(z.number()).default([]),
    prediction_duration: z.number().default(0),
    is_bargein: z.boolean().optional()
  }),
  z.object({
    type: z.literal(MSG_ERROR),
    message: z.string(),
    code: z.number().optional(),
    session_id: z.string().optional()
  })
]);
function resolveEffectiveThreshold(threshold, defaultThreshold) {
  if (threshold !== void 0) {
    return threshold;
  }
  if (defaultThreshold != null) {
    return defaultThreshold;
  }
  return null;
}
async function connectWebSocket(options) {
  const baseUrl = options.baseUrl.replace(/^http/, "ws");
  const token = await createAccessToken(options.apiKey, options.apiSecret);
  const url = `${baseUrl}/bargein`;
  const ws = new WebSocket(url, {
    headers: { ...buildMetadataHeaders(), Authorization: `Bearer ${token}` }
  });
  try {
    await new ThrowsPromise(
      (resolve, reject) => {
        const timeout = setTimeout(() => {
          ws.terminate();
          reject(
            new APITimeoutError({
              message: "WebSocket connection timeout",
              options: { retryable: false }
            })
          );
        }, options.connectTimeout);
        ws.once("open", () => {
          clearTimeout(timeout);
          resolve();
        });
        ws.once("unexpected-response", (_req, res) => {
          clearTimeout(timeout);
          ws.terminate();
          const statusCode = res.statusCode ?? -1;
          reject(
            new APIStatusError({
              message: `WebSocket connection rejected with status ${statusCode}`,
              options: { statusCode, retryable: false }
            })
          );
        });
        ws.once("error", (err) => {
          clearTimeout(timeout);
          ws.terminate();
          reject(new APIConnectionError({ message: `WebSocket connection error: ${err.message}` }));
        });
      }
    );
  } finally {
    ws.removeAllListeners();
  }
  return ws;
}
function createWsTransport(options, getState, setState, updateUserSpeakingSpan, onRequestSent, getAndResetNumRequests) {
  const logger = log();
  let activeWs = null;
  let outputController = null;
  const reconnectEvent = new Event();
  let reconnecting = null;
  let closed = false;
  function setupMessageHandler(socket) {
    socket.on("message", (data) => {
      let message;
      try {
        message = wsMessageSchema.parse(JSON.parse(data.toString()));
      } catch (err) {
        logger.warn(
          { data: data.toString(), err: err instanceof Error ? err.message : String(err) },
          "Failed to parse WebSocket message"
        );
        return;
      }
      try {
        handleMessage(message);
      } catch (err) {
        logger.warn(
          { type: message.type, err: err instanceof Error ? err.message : String(err) },
          "Failed to handle WebSocket message"
        );
      }
    });
    socket.on("error", (err) => {
      outputController == null ? void 0 : outputController.error(
        new APIConnectionError({ message: `WebSocket error: ${err.message}` })
      );
    });
    socket.on("close", (code, reason) => {
      logger.debug({ code, reason: reason.toString() }, "WebSocket closed");
    });
  }
  async function ensureConnection() {
    if (activeWs && activeWs.readyState === WebSocket.OPEN) return;
    activeWs = await connectWebSocket(options);
    setupMessageHandler(activeWs);
    const settings = {
      sample_rate: options.sampleRate,
      num_channels: 1,
      min_frames: options.minFrames,
      encoding: "s16le"
    };
    if (options.threshold !== void 0) {
      settings.threshold = options.threshold;
    }
    const sessionCreateMsg = JSON.stringify({
      type: MSG_SESSION_CREATE,
      settings
    });
    activeWs.send(sessionCreateMsg);
  }
  function handleMessage(message) {
    const state = getState();
    switch (message.type) {
      case MSG_SESSION_CREATED: {
        if (options.threshold === void 0 && message.default_threshold == null) {
          outputController == null ? void 0 : outputController.error(
            new APIStatusError({
              message: "adaptive interruption session created without a threshold: no user override and the server did not report a default_threshold",
              options: { statusCode: 500, retryable: false }
            })
          );
          break;
        }
        logger.debug(
          {
            defaultThreshold: message.default_threshold,
            effectiveThreshold: resolveEffectiveThreshold(
              options.threshold,
              message.default_threshold
            ),
            userOverride: options.threshold !== void 0
          },
          "adaptive interruption session created"
        );
        break;
      }
      case MSG_INTERRUPTION_DETECTED: {
        const createdAt = message.created_at;
        const overlapSpeechStartedAt = state.overlapSpeechStartedAt;
        if (state.overlapSpeechStarted && overlapSpeechStartedAt !== void 0) {
          const existing = state.cache.get(createdAt);
          const totalDurationInS = (existing == null ? void 0 : existing.requestStartedAt) !== void 0 ? (performance.now() - existing.requestStartedAt) / 1e3 : (performance.now() - createdAt) / 1e3;
          const entry = state.cache.setOrUpdate(
            createdAt,
            () => new InterruptionCacheEntry({ createdAt }),
            {
              speechInput: existing == null ? void 0 : existing.speechInput,
              requestStartedAt: existing == null ? void 0 : existing.requestStartedAt,
              totalDurationInS,
              probabilities: message.probabilities,
              isInterruption: true,
              predictionDurationInS: message.prediction_duration,
              detectionDelayInS: (Date.now() - overlapSpeechStartedAt) / 1e3
            }
          );
          if (updateUserSpeakingSpan) {
            updateUserSpeakingSpan(entry);
          }
          logger.debug(
            {
              totalDuration: entry.totalDurationInS,
              predictionDuration: entry.predictionDurationInS,
              detectionDelay: entry.detectionDelayInS,
              probability: entry.probability
            },
            "interruption detected"
          );
          const event = {
            type: "overlapping_speech",
            detectedAt: Date.now(),
            isInterruption: true,
            totalDurationInS: entry.totalDurationInS,
            predictionDurationInS: entry.predictionDurationInS,
            overlapStartedAt: overlapSpeechStartedAt,
            speechInput: entry.speechInput,
            probabilities: entry.probabilities,
            detectionDelayInS: entry.detectionDelayInS,
            probability: entry.probability,
            numRequests: (getAndResetNumRequests == null ? void 0 : getAndResetNumRequests()) ?? 0
          };
          if (outputController !== null && outputController.desiredSize !== null) {
            outputController.enqueue(event);
          } else {
            logger.debug("interruption output stream closed; dropping late bargein event");
          }
          setState({ overlapSpeechStarted: false });
        }
        break;
      }
      case MSG_INFERENCE_DONE: {
        const createdAt = message.created_at;
        const overlapSpeechStartedAt = state.overlapSpeechStartedAt;
        if (state.overlapSpeechStarted && overlapSpeechStartedAt !== void 0) {
          const existing = state.cache.get(createdAt);
          const totalDurationInS = (existing == null ? void 0 : existing.requestStartedAt) !== void 0 ? (performance.now() - existing.requestStartedAt) / 1e3 : (performance.now() - createdAt) / 1e3;
          const entry = state.cache.setOrUpdate(
            createdAt,
            () => new InterruptionCacheEntry({ createdAt }),
            {
              speechInput: existing == null ? void 0 : existing.speechInput,
              requestStartedAt: existing == null ? void 0 : existing.requestStartedAt,
              totalDurationInS,
              predictionDurationInS: message.prediction_duration,
              probabilities: message.probabilities,
              isInterruption: message.is_bargein ?? false,
              detectionDelayInS: (Date.now() - overlapSpeechStartedAt) / 1e3
            }
          );
          logger.debug(
            {
              totalDurationInS: entry.totalDurationInS,
              predictionDurationInS: entry.predictionDurationInS
            },
            "interruption inference done"
          );
        }
        break;
      }
      case MSG_SESSION_CLOSED:
        logger.debug("WebSocket session closed");
        break;
      case MSG_ERROR:
        outputController == null ? void 0 : outputController.error(
          new APIStatusError({
            message: `LiveKit Adaptive Interruption error: ${message.message}`,
            options: { statusCode: message.code ?? -1 }
          })
        );
        break;
    }
  }
  function sendAudioData(audioSlice) {
    if (!activeWs || activeWs.readyState !== WebSocket.OPEN) {
      throw new APIConnectionError({ message: "WebSocket not connected" });
    }
    const state = getState();
    const createdAt = Math.floor(performance.now());
    state.cache.set(
      createdAt,
      new InterruptionCacheEntry({
        createdAt,
        requestStartedAt: performance.now(),
        speechInput: audioSlice
      })
    );
    const header = new ArrayBuffer(8);
    const view = new DataView(header);
    view.setUint32(0, createdAt >>> 0, true);
    view.setUint32(4, Math.floor(createdAt / 4294967296) >>> 0, true);
    const audioBytes = new Uint8Array(
      audioSlice.buffer,
      audioSlice.byteOffset,
      audioSlice.byteLength
    );
    const combined = new Uint8Array(8 + audioBytes.length);
    combined.set(new Uint8Array(header), 0);
    combined.set(audioBytes, 8);
    activeWs.send(combined);
    onRequestSent == null ? void 0 : onRequestSent();
  }
  function closeSocket() {
    if ((activeWs == null ? void 0 : activeWs.readyState) === WebSocket.OPEN) {
      const closeMsg = JSON.stringify({ type: MSG_SESSION_CLOSE });
      try {
        activeWs.send(closeMsg);
      } catch (e) {
        logger.error(e, "failed to send close message");
      }
    }
    activeWs == null ? void 0 : activeWs.removeAllListeners();
    activeWs == null ? void 0 : activeWs.close(1e3);
    activeWs = null;
  }
  function close() {
    closed = true;
    reconnectEvent.set();
    closeSocket();
  }
  async function reconnect() {
    if (closed) return;
    reconnectEvent.set();
  }
  async function reconnectWatcher() {
    while (!closed) {
      await reconnectEvent.wait();
      if (closed) break;
      reconnectEvent.clear();
      const done = (async () => {
        closeSocket();
        getState().cache.clear();
        await ensureConnection();
        if (closed) closeSocket();
      })().catch((e) => {
        outputController == null ? void 0 : outputController.error(e);
      });
      reconnecting = done;
      try {
        await done;
      } finally {
        if (reconnecting === done) reconnecting = null;
      }
    }
  }
  const transport = new TransformStream(
    {
      async start(controller) {
        outputController = controller;
        await ensureConnection().catch((e) => {
          controller.error(e);
        });
        void reconnectWatcher();
      },
      async transform(chunk, controller) {
        if (!(chunk instanceof Int16Array)) {
          controller.enqueue(chunk);
          return;
        }
        if (reconnecting) await reconnecting;
        const state = getState();
        if (!state.overlapSpeechStartedAt || !state.overlapSpeechStarted) return;
        if (options.timeout > 0) {
          const now = performance.now();
          for (const [, entry] of state.cache.entries()) {
            if (entry.totalDurationInS !== 0) continue;
            if (now - entry.createdAt > options.timeout) {
              controller.error(
                new APIStatusError({
                  message: `interruption inference timed out after ${((now - entry.createdAt) / 1e3).toFixed(1)}s (ws)`,
                  options: { statusCode: 408, retryable: false }
                })
              );
              return;
            }
            break;
          }
        }
        try {
          sendAudioData(chunk);
        } catch (err) {
          controller.error(err);
        }
      },
      flush() {
        close();
      }
    },
    { highWaterMark: 2 },
    { highWaterMark: 2 }
  );
  return { transport, reconnect, close };
}
export {
  createWsTransport,
  resolveEffectiveThreshold,
  wsMessageSchema
};
//# sourceMappingURL=ws_transport.js.map