import assert from "node:assert/strict";
import test from "node:test";
import type {
  ExtensionAPI,
  ExtensionCommandContext,
  ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import {
  createDictateExtension,
  MAX_FRAME_BYTES,
  MAX_TRANSCRIPT_BYTES,
  PROTOCOL_VERSION,
  type DictateSocket,
} from "./index.ts";

type ShortcutHandler = (ctx: ExtensionContext) => Promise<void> | void;
type CommandHandler = (
  args: string,
  ctx: ExtensionCommandContext,
) => Promise<void> | void;
type EventHandler = (event: unknown, ctx: ExtensionContext) => Promise<void> | void;
type Frame = Record<string, unknown>;

class FakeSocket implements DictateSocket {
  readonly writes: string[] = [];
  endCount = 0;
  destroyCount = 0;
  private connectListener: (() => void) | undefined;
  private dataListener: ((chunk: Uint8Array) => void) | undefined;
  private errorListener: ((error: Error) => void) | undefined;
  private closeListener: (() => void) | undefined;

  onConnect(listener: () => void) {
    this.connectListener = listener;
  }
  onData(listener: (chunk: Uint8Array) => void) {
    this.dataListener = listener;
  }
  onError(listener: (error: Error) => void) {
    this.errorListener = listener;
  }
  onClose(listener: () => void) {
    this.closeListener = listener;
  }
  write(data: string) {
    this.writes.push(data);
  }
  end() {
    this.endCount++;
  }
  destroy() {
    this.destroyCount++;
  }
  connect() {
    this.connectListener?.();
  }
  data(data: string | Uint8Array) {
    this.dataListener?.(
      typeof data === "string" ? Buffer.from(data, "utf8") : data,
    );
  }
  frame(frame: Frame) {
    this.data(`${JSON.stringify(frame)}\n`);
  }
  error(message: string) {
    this.errorListener?.(new Error(message));
  }
  close() {
    this.closeListener?.();
  }
}

function createHarness() {
  const shortcuts = new Map<string, ShortcutHandler>();
  const commands = new Map<string, CommandHandler>();
  const events = new Map<string, EventHandler>();
  const sockets: FakeSocket[] = [];
  const socketPaths: string[] = [];
  const statuses: Array<string | undefined> = [];
  const themeCalls: Array<{ color: string; text: string }> = [];
  const notifications: Array<{ message: string; type?: string }> = [];
  const pasted: string[] = [];
  const context = {
    ui: {
      theme: {
        fg: (color: string, text: string) => {
          themeCalls.push({ color, text });
          return text;
        },
      },
      setStatus: (_key: string, text: string | undefined) => statuses.push(text),
      notify: (message: string, type?: string) =>
        notifications.push({ message, type }),
      pasteToEditor: (text: string) => pasted.push(text),
    },
  } as unknown as ExtensionContext;
  const api = {
    registerShortcut(key: string, options: { handler: ShortcutHandler }) {
      shortcuts.set(key, options.handler);
    },
    registerCommand(name: string, options: { handler: CommandHandler }) {
      commands.set(name, options.handler);
    },
    on(name: string, handler: EventHandler) {
      events.set(name, handler);
    },
  } as unknown as ExtensionAPI;

  createDictateExtension({
    socketPath: "/test/pi-dictate.sock",
    socketFactory: (path) => {
      socketPaths.push(path);
      const socket = new FakeSocket();
      sockets.push(socket);
      return socket;
    },
  })(api);

  const shortcut = (key: string) => {
    const handler = shortcuts.get(key);
    assert.ok(handler, `shortcut ${key} was registered`);
    return handler;
  };
  const command = (name: string) => {
    const handler = commands.get(name);
    assert.ok(handler, `command ${name} was registered`);
    return handler;
  };
  const event = (name: string) => {
    const handler = events.get(name);
    assert.ok(handler, `event ${name} was registered`);
    return handler;
  };
  const sentFrames = (socket: FakeSocket) =>
    socket.writes.map((line) => JSON.parse(line) as Frame);

  return {
    shortcuts,
    commands,
    events,
    sockets,
    socketPaths,
    statuses,
    themeCalls,
    notifications,
    pasted,
    context,
    shortcut,
    command,
    event,
    sentFrames,
  };
}

type Harness = ReturnType<typeof createHarness>;

async function connectAndNegotiate(harness: Harness) {
  await harness.shortcut("alt+m")(harness.context);
  const socket = harness.sockets.at(-1);
  assert.ok(socket);
  socket.connect();
  const hello = harness.sentFrames(socket)[0];
  assert.equal(hello?.type, "hello");
  assert.equal(hello?.version, PROTOCOL_VERSION);
  assert.equal(typeof hello?.requestId, "string");
  assert.equal(harness.sentFrames(socket).length, 1, "start waits for hello");
  socket.frame({
    version: PROTOCOL_VERSION,
    type: "pong",
    requestId: hello?.requestId,
  });
  const start = harness.sentFrames(socket)[1];
  assert.equal(start?.type, "start");
  assert.equal(start?.version, PROTOCOL_VERSION);
  assert.equal(typeof start?.requestId, "string");
  assert.equal(typeof start?.generationId, "string");
  return { socket, generationId: start?.generationId as string };
}

function serviceFrame(
  socket: FakeSocket,
  generationId: string,
  type: string,
  fields: Frame = {},
) {
  socket.frame({
    version: PROTOCOL_VERSION,
    type,
    requestId: `event-${type}`,
    generationId,
    ...fields,
  });
}

test("registers main-editor shortcuts, cancel command, and shutdown cleanup", () => {
  const harness = createHarness();
  assert.deepEqual([...harness.shortcuts.keys()], ["alt+m", "alt+n"]);
  assert.deepEqual([...harness.commands.keys()], ["dictate-cancel", "dictate-status"]);
  assert.ok(harness.events.has("session_shutdown"));
});

test("status command negotiates a short read-only connection and reports backend diagnostics", async () => {
  const harness = createHarness();
  await harness.command("dictate-status")(
    "",
    harness.context as ExtensionCommandContext,
  );
  const socket = harness.sockets[0];
  assert.ok(socket);
  socket.connect();
  const hello = harness.sentFrames(socket)[0];
  assert.equal(hello?.type, "hello");
  socket.frame({
    version: PROTOCOL_VERSION,
    type: "pong",
    requestId: hello?.requestId,
    backend: "openai",
    model: "gpt-live-transcribe",
  });
  const ping = harness.sentFrames(socket)[1];
  assert.equal(ping?.type, "ping");
  socket.frame({
    version: PROTOCOL_VERSION,
    type: "pong",
    requestId: ping?.requestId,
  });

  assert.deepEqual(harness.notifications, [
    {
      message: "Dictation backend: openai; model: gpt-live-transcribe",
      type: "info",
    },
  ]);
  assert.equal(socket.endCount, 1);
});

test("status diagnostics reject control characters without exposing them", async () => {
  const harness = createHarness();
  await harness.command("dictate-status")(
    "",
    harness.context as ExtensionCommandContext,
  );
  const socket = harness.sockets[0];
  assert.ok(socket);
  socket.connect();
  const hello = harness.sentFrames(socket)[0];
  socket.frame({
    version: PROTOCOL_VERSION,
    type: "pong",
    requestId: hello?.requestId,
    backend: "openai\nsecret",
  });

  assert.match(harness.notifications[0]?.message ?? "", /diagnostic backend is invalid/i);
  assert.equal(socket.destroyCount, 1);
});

test("connection errors return to idle and the next start uses a fresh socket", async () => {
  const harness = createHarness();
  await harness.shortcut("alt+m")(harness.context);
  harness.sockets[0]?.error("socket unavailable");

  assert.match(harness.notifications[0]?.message ?? "", /connection failed.*socket unavailable/i);
  assert.equal(harness.notifications[0]?.type, "error");
  assert.equal(harness.statuses.at(-1), undefined);
  assert.equal(harness.sockets[0]?.destroyCount, 1);

  await harness.shortcut("alt+m")(harness.context);
  assert.equal(harness.sockets.length, 2);
  assert.deepEqual(harness.socketPaths, ["/test/pi-dictate.sock", "/test/pi-dictate.sock"]);
});

test("rejects a mismatched protocol version before sending start", async () => {
  const harness = createHarness();
  await harness.shortcut("alt+m")(harness.context);
  const socket = harness.sockets[0];
  assert.ok(socket);
  socket.connect();
  const hello = harness.sentFrames(socket)[0];
  socket.frame({
    version: PROTOCOL_VERSION + 1,
    type: "pong",
    requestId: hello?.requestId,
  });

  assert.equal(harness.sentFrames(socket).length, 1);
  assert.match(harness.notifications[0]?.message ?? "", /version mismatch/i);
  assert.equal(socket.destroyCount, 1);
});

test("rejects a hello response with the wrong request ID", async () => {
  const harness = createHarness();
  await harness.shortcut("alt+m")(harness.context);
  const socket = harness.sockets[0];
  assert.ok(socket);
  socket.connect();
  socket.frame({
    version: PROTOCOL_VERSION,
    type: "pong",
    requestId: "not-the-hello-request",
  });

  assert.equal(harness.sentFrames(socket).length, 1);
  assert.match(harness.notifications[0]?.message ?? "", /unexpected hello response/i);
  assert.equal(socket.destroyCount, 1);
});

test("rejects generation events before hello negotiation", async () => {
  const harness = createHarness();
  await harness.shortcut("alt+m")(harness.context);
  const socket = harness.sockets[0];
  assert.ok(socket);
  socket.connect();
  serviceFrame(socket, "generation-1", "state", { state: "listening" });

  assert.match(harness.notifications[0]?.message ?? "", /before hello negotiation/i);
  assert.equal(socket.destroyCount, 1);
  assert.deepEqual(harness.pasted, []);
});

test("renders authoritative service states and sends a bounded stop command", async () => {
  const harness = createHarness();
  const { socket, generationId } = await connectAndNegotiate(harness);

  serviceFrame(socket, generationId, "state", { state: "loading" });
  serviceFrame(socket, generationId, "state", { state: "listening" });
  await harness.shortcut("alt+m")(harness.context);
  const stop = harness.sentFrames(socket).at(-1);
  assert.equal(stop?.type, "stop");
  assert.equal(stop?.generationId, generationId);
  assert.ok(Buffer.byteLength(String(stop?.requestId)) <= 128);
  serviceFrame(socket, generationId, "state", { state: "stopping" });
  serviceFrame(socket, generationId, "state", { state: "transcribing" });

  assert.deepEqual(harness.statuses, [
    "preparing transcription...",
    "● listening  ▁▁▁▁▁▁▁▁▁▁▁▁",
    "stopping...",
    "transcribing...",
  ]);
});

test("scrolls bounded audio levels through the listening waveform", async () => {
  const harness = createHarness();
  const { socket, generationId } = await connectAndNegotiate(harness);
  serviceFrame(socket, generationId, "state", { state: "listening" });
  assert.deepEqual(harness.themeCalls, [
    { color: "error", text: "●" },
    { color: "dim", text: "listening" },
    { color: "accent", text: "▁▁▁▁▁▁▁▁▁▁▁▁" },
  ]);
  serviceFrame(socket, generationId, "level", { level: 0.5 });
  serviceFrame(socket, generationId, "level", { level: 1 });

  assert.deepEqual(harness.statuses, [
    "● listening  ▁▁▁▁▁▁▁▁▁▁▁▁",
    "● listening  ▁▁▁▁▁▁▁▁▁▁▁▅",
    "● listening  ▁▁▁▁▁▁▁▁▁▁▅█",
  ]);
});

test("rejects audio levels outside listening or the normalized range", async () => {
  const beforeListening = createHarness();
  const pending = await connectAndNegotiate(beforeListening);
  serviceFrame(pending.socket, pending.generationId, "level", { level: 0.5 });
  assert.match(beforeListening.notifications[0]?.message ?? "", /audio level is invalid/i);

  const outOfRange = createHarness();
  const active = await connectAndNegotiate(outOfRange);
  serviceFrame(active.socket, active.generationId, "state", { state: "listening" });
  serviceFrame(active.socket, active.generationId, "level", { level: 1.1 });
  assert.match(outOfRange.notifications[0]?.message ?? "", /audio level is invalid/i);
});

test("ignores stale generations but rejects malformed and oversized frames", async () => {
  const staleHarness = createHarness();
  const { socket: staleSocket, generationId } = await connectAndNegotiate(staleHarness);
  serviceFrame(staleSocket, "old-generation", "state", { state: "listening" });
  assert.deepEqual(staleHarness.statuses, []);
  serviceFrame(staleSocket, generationId, "state", { state: "listening" });
  assert.deepEqual(staleHarness.statuses, ["● listening  ▁▁▁▁▁▁▁▁▁▁▁▁"]);

  const malformedHarness = createHarness();
  const malformed = await connectAndNegotiate(malformedHarness);
  malformed.socket.data("{not json}\n");
  assert.match(malformedHarness.notifications[0]?.message ?? "", /malformed JSON/i);
  assert.equal(malformed.socket.destroyCount, 1);

  const nonObjectHarness = createHarness();
  const nonObject = await connectAndNegotiate(nonObjectHarness);
  nonObject.socket.data("null\n");
  assert.match(nonObjectHarness.notifications[0]?.message ?? "", /frame must be an object/i);
  assert.equal(nonObject.socket.destroyCount, 1);

  const oversizedHarness = createHarness();
  const oversized = await connectAndNegotiate(oversizedHarness);
  oversized.socket.data("x".repeat(MAX_FRAME_BYTES + 1));
  assert.match(oversizedHarness.notifications[0]?.message ?? "", /frame is too large/i);
  assert.equal(oversized.socket.destroyCount, 1);
});

test("rejects invalid UTF-8 instead of altering transcript text", async () => {
  const harness = createHarness();
  const { socket, generationId } = await connectAndNegotiate(harness);
  serviceFrame(socket, generationId, "state", { state: "listening" });
  await harness.shortcut("alt+m")(harness.context);
  serviceFrame(socket, generationId, "state", { state: "transcribing" });
  const prefix = Buffer.from(
    JSON.stringify({
      version: PROTOCOL_VERSION,
      type: "result",
      requestId: "event-result",
      generationId,
      transcript: "",
    }).replace('"transcript":""}', '"transcript":"'),
  );
  const suffix = Buffer.from('"}\n');
  socket.data(Buffer.concat([prefix, Buffer.from([0xc3, 0x28]), suffix]));

  assert.match(harness.notifications[0]?.message ?? "", /invalid UTF-8/i);
  assert.equal(socket.destroyCount, 1);
  assert.deepEqual(harness.pasted, []);
});

test("service idle closes the accepted connection and forces a fresh handshake", async () => {
  const harness = createHarness();
  const { socket, generationId } = await connectAndNegotiate(harness);
  serviceFrame(socket, generationId, "state", { state: "loading" });
  serviceFrame(socket, generationId, "state", { state: "idle" });
  assert.equal(socket.endCount, 1);
  assert.equal(harness.statuses.at(-1), undefined);

  await harness.shortcut("alt+m")(harness.context);
  assert.equal(harness.sockets.length, 2);
  harness.sockets[1]?.connect();
  assert.equal(harness.sentFrames(harness.sockets[1] as FakeSocket)[0]?.type, "hello");
});

test("an empty result leaves the editor unchanged and reports no speech", async () => {
  const harness = createHarness();
  const { socket, generationId } = await connectAndNegotiate(harness);
  serviceFrame(socket, generationId, "state", { state: "listening" });
  await harness.shortcut("alt+m")(harness.context);
  serviceFrame(socket, generationId, "state", { state: "transcribing" });
  serviceFrame(socket, generationId, "result", { transcript: "  " });

  assert.deepEqual(harness.pasted, []);
  assert.match(harness.notifications[0]?.message ?? "", /no speech recognized/i);
  assert.equal(socket.endCount, 1);
});

test("pastes exactly one non-empty current result through the public editor API", async () => {
  const harness = createHarness();
  const { socket, generationId } = await connectAndNegotiate(harness);
  serviceFrame(socket, generationId, "state", { state: "listening" });
  await harness.shortcut("alt+m")(harness.context);
  serviceFrame(socket, generationId, "state", { state: "stopping" });
  serviceFrame(socket, generationId, "state", { state: "transcribing" });
  serviceFrame(socket, generationId, "result", { transcript: "typed words" });
  serviceFrame(socket, generationId, "result", { transcript: "duplicate" });

  assert.deepEqual(harness.pasted, ["typed words"]);
  assert.equal(harness.statuses.at(-1), undefined);

  const oversizedTranscriptHarness = createHarness();
  const active = await connectAndNegotiate(oversizedTranscriptHarness);
  serviceFrame(active.socket, active.generationId, "state", { state: "listening" });
  await oversizedTranscriptHarness.shortcut("alt+m")(oversizedTranscriptHarness.context);
  serviceFrame(active.socket, active.generationId, "state", { state: "transcribing" });
  serviceFrame(active.socket, active.generationId, "result", {
    transcript: "t".repeat(MAX_TRANSCRIPT_BYTES + 1),
  });
  assert.deepEqual(oversizedTranscriptHarness.pasted, []);
  assert.match(oversizedTranscriptHarness.notifications[0]?.message ?? "", /transcript is too large/i);
});

test("disconnects and service errors return safely to idle without delivery", async () => {
  const disconnected = createHarness();
  const first = await connectAndNegotiate(disconnected);
  first.socket.close();
  assert.match(disconnected.notifications[0]?.message ?? "", /disconnected/i);
  assert.deepEqual(disconnected.pasted, []);
  await disconnected.shortcut("alt+m")(disconnected.context);
  assert.equal(disconnected.sockets.length, 2);

  const failed = createHarness();
  const second = await connectAndNegotiate(failed);
  serviceFrame(second.socket, second.generationId, "error", {
    code: "model_failed",
    detail: "model unavailable",
  });
  assert.match(failed.notifications[0]?.message ?? "", /model_failed.*model unavailable/i);
  assert.deepEqual(failed.pasted, []);
});

test("cancellation during loading tombstones later events and the slash command also cancels", async () => {
  const harness = createHarness();
  const { socket, generationId } = await connectAndNegotiate(harness);
  serviceFrame(socket, generationId, "state", { state: "loading" });
  await harness.shortcut("alt+n")(harness.context);

  assert.equal(harness.sentFrames(socket).at(-1)?.type, "cancel");
  assert.equal(socket.endCount, 1);
  serviceFrame(socket, generationId, "state", { state: "listening" });
  serviceFrame(socket, generationId, "result", { transcript: "must not appear" });
  assert.deepEqual(harness.pasted, []);
  assert.equal(harness.statuses.at(-1), undefined);

  const next = await connectAndNegotiate(harness);
  serviceFrame(next.socket, next.generationId, "state", { state: "listening" });
  await harness.command("dictate-cancel")(
    "",
    harness.context as ExtensionCommandContext,
  );
  assert.equal(harness.sentFrames(next.socket).at(-1)?.type, "cancel");
  assert.equal(next.socket.endCount, 1);
});

test("shutdown best-effort cancels, closes, clears status, and remains delivery-neutral", async () => {
  const harness = createHarness();
  const { socket, generationId } = await connectAndNegotiate(harness);
  serviceFrame(socket, generationId, "state", { state: "listening" });
  const shutdown = harness.event("session_shutdown");
  await shutdown({}, harness.context);
  await shutdown({}, harness.context);

  assert.equal(harness.sentFrames(socket).at(-1)?.type, "cancel");
  assert.equal(socket.endCount, 1);
  assert.equal(harness.statuses.at(-1), undefined);
  serviceFrame(socket, generationId, "result", { transcript: "never paste cleanup" });
  assert.deepEqual(harness.pasted, []);
});