import assert from "node:assert/strict";
import { afterEach, test } from "node:test";

import { VoiceSessionController } from
  "../../src/voice_assistant/audio_gateway/web/voice-session.js";

const originalWebSocket = globalThis.WebSocket;

class FakeWebSocket {
  static OPEN = 1;
  static CLOSED = 3;

  constructor() {
    this.readyState = FakeWebSocket.OPEN;
    this.sent = [];
    this.closed = false;
  }

  send(data) {
    this.sent.push(data);
  }

  close() {
    this.closed = true;
    this.readyState = FakeWebSocket.CLOSED;
  }
}

globalThis.WebSocket = FakeWebSocket;
globalThis.cancelAnimationFrame = () => {};
globalThis.requestAnimationFrame = () => 1;


afterEach(() => {
  // Restore only state modified by individual tests.
});

function makeController() {
  const states = [];
  const errors = [];
  const workletMessages = [];

  const controller = new VoiceSessionController({
    onStateChange: (state) => states.push(state),
    onError: (error) => errors.push(error),
  });

  controller.ws = new FakeWebSocket();
  controller.worklet = {
    port: {
      postMessage: (data) => workletMessages.push(data),
      onmessage: null,
    },
    disconnect() {},
  };

  controller.source = { disconnect() {} };
  controller.silentGain = { disconnect() {} };

  return {
    controller,
    states,
    errors,
    workletMessages,
  };
}

test("speech detection requests a new turn", () => {
  const { controller } = makeController();

  controller.setState("listening");
  controller.beginTurn();

  assert.equal(controller.state, "starting");

  assert.deepEqual(
    JSON.parse(controller.ws.sent[0]),
    { type: "audio.start" }
  );
});

test("PCM is buffered until turn.started", () => {
  const { controller } = makeController();

  controller.setState("listening");

  controller.handlePCM(new ArrayBuffer(4));

  assert.equal(controller.prebuffer.length, 1);
  assert.equal(controller.ws.sent.length, 0);

  controller.beginTurn();
  controller.handlePCM(new ArrayBuffer(6));

  assert.equal(controller.pending.length, 2);
  assert.equal(controller.ws.sent.length, 1);

  controller.handleGatewayEvent({
    type: "turn.started",
    turn_id: 1,
  });

  assert.equal(controller.state, "recording");
  assert.equal(controller.pending.length, 0);
  assert.equal(controller.sequence, 2);
  assert.equal(controller.ws.sent.length, 5);
});

test("PCM metadata precedes each binary frame", () => {
  const { controller } = makeController();

  controller.turnId = 7;
  controller.setState("recording");

  const pcm = new ArrayBuffer(8);

  controller.sendPCM(pcm);

  assert.deepEqual(
    JSON.parse(controller.ws.sent[0]),
    {
      type: "audio.chunk",
      turn_id: 7,
      sequence_number: 0,
    }
  );

  assert.equal(controller.ws.sent[1], pcm);
  assert.equal(controller.sequence, 1);
});

test("speech end requests worklet flush", () => {
  const { controller, workletMessages } = makeController();

  controller.turnId = 1;
  controller.setState("recording");

  controller.finishSpeech();

  assert.equal(controller.state, "flushing");
  assert.deepEqual(
    workletMessages.at(-1),
    { type: "flush" }
  );

  assert.equal(controller.ws.sent.length, 0);
});


test("audio.stop is sent after flushed", () => {
  const { controller } = makeController();

  controller.turnId = 1;
  controller.setState("flushing");

  controller.handleFlushed();

  const messages = controller.ws.sent
    .filter((message) => typeof message === "string")
    .map((message) => JSON.parse(message));

  assert.ok(
    messages.some(
      (message) =>
        message.type === "audio.stop" &&
        message.turn_id === 1
    )
  );

  assert.equal(controller.state, "processing");
});


test("speech ending before turn.started is handled", () => {
  const { controller, workletMessages } = makeController();

  controller.setState("listening");
  controller.beginTurn();
  controller.finishSpeech();

  assert.equal(controller.stopRequested, true);
  assert.equal(controller.state, "starting");

  controller.handleGatewayEvent({
    type: "turn.started",
    turn_id: 3,
  });

  assert.equal(controller.state, "flushing");
  assert.deepEqual(
    workletMessages.at(-1),
    { type: "flush" }
  );
});

test("closing a session releases resources", () => {
  const { controller, workletMessages } = makeController();
  const socket = controller.ws;

  controller.setState("listening");
  controller.close();

  assert.equal(controller.state, "closed");
  assert.equal(socket.closed, true);
  assert.equal(controller.ws, null);
  assert.deepEqual(
    workletMessages.at(-1),
    { type: "stop" }
  );
});

test("gateway errors close the session", () => {
  const { controller, errors } = makeController();

  controller.setState("processing");

  controller.handleGatewayEvent({
    type: "error",
    code: "STT_START_FAILED",
    message: "Provider unavailable",
  });

  assert.equal(controller.state, "closed");
  assert.equal(errors.length, 1);
  assert.match(
    errors[0].message,
    /STT_START_FAILED/
  );
});


test("turn.ended waits for playback before listening", async () => {
  const { controller } = makeController();

  let resolvePlayback;

  controller.player = {
    finish() {},
    reset() {},
    stop() {},
    waitForCompletion() {
      return new Promise((resolve) => {
        resolvePlayback = resolve;
      });
    },
  };

  controller.turnId = 1;
  controller.setState("processing");

  let resumed = false;

  controller.startListening = () => {
    resumed = true;
  };

  controller.handleGatewayEvent({
    type: "turn.ended",
    turn_id: 1,
  });

  await Promise.resolve();
  await Promise.resolve();

  assert.equal(resumed, false);

  resolvePlayback();

  await Promise.resolve();
  await Promise.resolve();

  assert.equal(resumed, true);
});

test("closing during playback prevents listening resume", async () => {
  const { controller } = makeController();

  let resolvePlayback;

  controller.player = {
    finish() {},
    reset() {},
    stop() {},
    waitForCompletion() {
      return new Promise((resolve) => {
        resolvePlayback = resolve;
      });
    },
  };

  controller.turnId = 1;
  controller.setState("processing");

  let resumed = false;

  controller.startListening = () => {
    resumed = true;
  };

  controller.handleGatewayEvent({
    type: "turn.ended",
    turn_id: 1,
  });

  await Promise.resolve();
  await Promise.resolve();

  controller.close();

  resolvePlayback();

  await Promise.resolve();
  await Promise.resolve();

  assert.equal(resumed, false);
  assert.equal(controller.state, "closed");
});

test("binary audio is forwarded to PCMPlayer", async () => {
  const { controller } = makeController();

  const received = [];

  controller.player = {
    async enqueue(buffer) {
      received.push(buffer);
    },
    finish() {},
    reset() {},
    stop() {},
    waitForCompletion: async () => {},
  };

  controller.turnId = 1;
  controller.setState("processing");

  const pcm = new ArrayBuffer(8);

  controller.handleAudioChunk(pcm);

  await controller.playbackPending;

  assert.deepEqual(received, [pcm]);
});



test("barge-in stops playback and starts a new turn", () => {
  const { controller } = makeController();

  let playbackStopped = false;

  controller.player = {
    stop() {
      playbackStopped = true;
    },
    reset() {},
  };

  controller.turnId = 1;
  controller.setState("speaking");

  // Simulate user speech detected during TTS.
  controller.requestBargeIn();

  assert.equal(playbackStopped, true);
  assert.equal(controller.state, "interrupting");
  assert.equal(controller.bargeInPending, true);

  assert.deepEqual(
    JSON.parse(controller.ws.sent[0]),
    {
      type: "interrupt",
      turn_id: 1,
    }
  );

  // Gateway confirms the previous turn was cancelled.
  controller.handleGatewayEvent({
    type: "interrupted",
    turn_id: 1,
  });

  assert.equal(controller.state, "starting");
  assert.equal(controller.ws.closed, false);

  assert.deepEqual(
    JSON.parse(controller.ws.sent[1]),
    {
      type: "audio.start",
    }
  );

  // Gateway accepts the new turn.
  controller.handleGatewayEvent({
    type: "turn.started",
    turn_id: 2,
  });

  assert.equal(controller.turnId, 2);
  assert.equal(controller.state, "recording");
  assert.equal(controller.bargeInPending, false);
});



test("barge-in preserves PCM until new turn starts", () => {
  const { controller } = makeController();

  controller.player = {
    stop() {},
    reset() {},
  };

  controller.turnId = 1;
  controller.setState("speaking");

  const pcm1 = new ArrayBuffer(4);
  const pcm2 = new ArrayBuffer(6);
  const pcm3 = new ArrayBuffer(8);

  // Audio captured before speech confirmation.
  controller.handlePCM(pcm1);

  controller.requestBargeIn();

  // Audio captured while interruption is pending.
  controller.handlePCM(pcm2);

  assert.equal(controller.ws.sent.length, 1);
  assert.equal(controller.bargeInBuffer.length, 2);

  controller.handleGatewayEvent({
    type: "interrupted",
    turn_id: 1,
  });

  // Still no PCM sent to the old turn.
  assert.equal(controller.state, "starting");
  assert.equal(controller.ws.sent.length, 2);

  // Additional speech arrives before turn.started.
  controller.handlePCM(pcm3);

  controller.handleGatewayEvent({
    type: "turn.started",
    turn_id: 2,
  });

  const sent = controller.ws.sent;

  assert.deepEqual(
    JSON.parse(sent[0]),
    { type: "interrupt", turn_id: 1 }
  );

  assert.deepEqual(
    JSON.parse(sent[1]),
    { type: "audio.start" }
  );

  assert.deepEqual(
    JSON.parse(sent[2]),
    {
      type: "audio.chunk",
      turn_id: 2,
      sequence_number: 0,
    }
  );

  assert.equal(sent[3], pcm1);

  assert.deepEqual(
    JSON.parse(sent[4]),
    {
      type: "audio.chunk",
      turn_id: 2,
      sequence_number: 1,
    }
  );

  assert.equal(sent[5], pcm2);

  assert.deepEqual(
    JSON.parse(sent[6]),
    {
      type: "audio.chunk",
      turn_id: 2,
      sequence_number: 2,
    }
  );

  assert.equal(sent[7], pcm3);

  assert.equal(controller.sequence, 3);
  assert.equal(controller.state, "recording");
});


test("barge-in finishes speech before new turn starts", () => {
  const { controller, workletMessages } = makeController();

  controller.player = {
    stop() {},
    reset() {},
  };

  controller.turnId = 1;
  controller.setState("speaking");

  const pcm = new ArrayBuffer(8);

  controller.handlePCM(pcm);
  controller.requestBargeIn();

  // User finishes speaking before the Gateway responds.
  controller.finishBargeInSpeech();

  assert.equal(controller.bargeInSpeechEnded, true);

  controller.handleGatewayEvent({
    type: "interrupted",
    turn_id: 1,
  });

  assert.equal(controller.state, "starting");

  controller.handleGatewayEvent({
    type: "turn.started",
    turn_id: 2,
  });

  // Buffered audio must be transmitted before flushing.
  assert.equal(controller.sequence, 1);

  assert.deepEqual(
    JSON.parse(controller.ws.sent[2]),
    {
      type: "audio.chunk",
      turn_id: 2,
      sequence_number: 0,
    }
  );

  assert.equal(controller.ws.sent[3], pcm);
  assert.equal(controller.state, "flushing");

  assert.deepEqual(
    workletMessages.at(-1),
    { type: "flush" }
  );
});




test("barge-in continues recording until speech ends", () => {
  const { controller, workletMessages } = makeController();

  controller.player = {
    stop() {},
    reset() {},
  };

  controller.turnId = 1;
  controller.setState("speaking");

  controller.bargeInVad.start();
  controller.requestBargeIn();

  controller.handleGatewayEvent({
    type: "interrupted",
    turn_id: 1,
  });

  controller.handleGatewayEvent({
    type: "turn.started",
    turn_id: 2,
  });

  assert.equal(controller.state, "recording");
  assert.equal(controller.bargeInRecording, true);
  assert.equal(controller.bargeInVad.active, true);

  // Speech continues after the new turn starts.
  const pcm = new ArrayBuffer(8);
  controller.handlePCM(pcm);

  assert.equal(controller.sequence, 1);

  // Barge-in VAD detects the end of speech.
  controller.finishBargeInSpeech();

  assert.equal(controller.bargeInRecording, false);
  assert.equal(controller.state, "flushing");

  assert.deepEqual(
    workletMessages.at(-1),
    { type: "flush" }
  );

  // Worklet confirms that remaining PCM was emitted.
  controller.handleFlushed();

  assert.equal(controller.state, "processing");

  const messages = controller.ws.sent
    .filter((message) => typeof message === "string")
    .map((message) => JSON.parse(message));

  assert.ok(
    messages.some(
      (message) =>
        message.type === "audio.stop" &&
        message.turn_id === 2
    )
  );
});


test("late INVALID_TURN recovers barge-in session", () => {
  const { controller } = makeController();

  controller.player = {
    stop() {},
    reset() {},
  };

  controller.turnId = 1;
  controller.setState("speaking");

  controller.requestBargeIn();

  assert.equal(controller.state, "interrupting");

  controller.handleGatewayEvent({
    type: "error",
    code: "INVALID_TURN",
    message: "Turn 1 is no longer active.",
    turn_id: 1,
  });

  assert.equal(controller.state, "starting");
  assert.equal(controller.ws.closed, false);

  const messages = controller.ws.sent
    .filter((message) => typeof message === "string")
    .map((message) => JSON.parse(message));

  assert.deepEqual(
    messages.map((message) => message.type),
    ["interrupt", "audio.start"]
  );

  controller.handleGatewayEvent({
    type: "turn.started",
    turn_id: 2,
  });

  assert.equal(controller.state, "recording");
  assert.equal(controller.turnId, 2);
});


test("late TTS chunk during barge-in does not close session", () => {
  const { controller } = makeController();

  let stopped = false;

  controller.player = {
    stop() {
      stopped = true;
    },
    reset() {},
  };

  controller.turnId = 1;
  controller.setState("speaking");

  controller.requestBargeIn();

  assert.equal(stopped, true);
  assert.equal(controller.state, "interrupting");

  // A late PCM chunk from the interrupted TTS.
  controller.handleAudioChunk(new ArrayBuffer(8));

  // The WebSocket session must remain alive.
  assert.equal(controller.ws.closed, false);
  assert.equal(controller.state, "interrupting");
  assert.equal(controller.bargeInPending, true);
});


test("barge-in resets player before new TTS response", () => {
  const { controller } = makeController();

  let finished = false;
  let resetCount = 0;

  controller.player = {
    stop() {
      finished = true;
    },

    reset() {
      finished = false;
      resetCount++;
    },
  };

  controller.turnId = 1;
  controller.setState("speaking");

  controller.requestBargeIn();

  assert.equal(finished, true);

  controller.handleGatewayEvent({
    type: "interrupted",
    turn_id: 1,
  });

  assert.equal(controller.state, "starting");

  controller.handleGatewayEvent({
    type: "turn.started",
    turn_id: 2,
  });

  assert.equal(controller.state, "recording");
  assert.equal(finished, false);
  assert.equal(resetCount, 1);
});