import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { createConnection } from "node:net";

export const DICTATE_SOCKET_PATH = "/run/user/1000/pi-dictate.sock";
export const PROTOCOL_VERSION = 2;
export const MAX_FRAME_BYTES = 64 * 1024;
export const MAX_TRANSCRIPT_BYTES = 32 * 1024;
export const MAX_ID_BYTES = 128;

const STATUS_KEY = "dictate";
const MAX_ERROR_DETAIL_BYTES = 2 * 1024;
const MAX_DIAGNOSTIC_BYTES = 128;
const UTF8_DECODER = new TextDecoder("utf-8", { fatal: true });
const WAVEFORM_BARS = ["▁", "▂", "▃", "▄", "▅", "▆", "▇", "█"] as const;
const WAVEFORM_SAMPLES = 12;

type DictateState =
  | "idle"
  | "connecting"
  | "recording"
  | "stopping"
  | "transcribing";
type ServiceState =
  | "loading"
  | "listening"
  | "stopping"
  | "transcribing"
  | "idle";

export interface DictateSocket {
  onConnect(listener: () => void): void;
  onData(listener: (chunk: Uint8Array) => void): void;
  onError(listener: (error: Error) => void): void;
  onClose(listener: () => void): void;
  write(data: string): void;
  end(): void;
  destroy(): void;
}

export interface DictateExtensionOptions {
  socketFactory?: (path: string) => DictateSocket;
  socketPath?: string;
}

function nodeSocket(path: string): DictateSocket {
  const socket = createConnection(path);
  return {
    onConnect: (listener) => socket.on("connect", listener),
    onData: (listener) => socket.on("data", listener),
    onError: (listener) => socket.on("error", listener),
    onClose: (listener) => socket.on("close", listener),
    write: (data) => {
      socket.write(data);
    },
    end: () => {
      socket.end();
    },
    destroy: () => {
      socket.destroy();
    },
  };
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === "object" && value !== null && !Array.isArray(value);
}

function utf8Bytes(value: string): number {
  return Buffer.byteLength(value, "utf8");
}

function isBoundedId(value: unknown): value is string {
  return (
    typeof value === "string" &&
    value.length > 0 &&
    utf8Bytes(value) <= MAX_ID_BYTES
  );
}

function isBoundedDiagnostic(value: unknown): value is string {
  return (
    typeof value === "string" &&
    value.length > 0 &&
    utf8Bytes(value) <= MAX_DIAGNOSTIC_BYTES &&
    !/[\u0000-\u001f\u007f]/u.test(value)
  );
}

function boundedText(value: string, maxBytes: number): string {
  if (utf8Bytes(value) <= maxBytes) return value;
  let text = value;
  while (text.length > 0 && utf8Bytes(`${text}…`) > maxBytes) {
    text = text.slice(0, -1);
  }
  return `${text}…`;
}

export function createDictateExtension(options: DictateExtensionOptions = {}) {
  const socketFactory = options.socketFactory ?? nodeSocket;
  const socketPath = options.socketPath ?? DICTATE_SOCKET_PATH;

  return function registerDictate(pi: ExtensionAPI) {
    let state: DictateState = "idle";
    let socket: DictateSocket | null = null;
    let statusSocket: DictateSocket | null = null;
    let activeContext: ExtensionContext | null = null;
    let statusContext: ExtensionContext | null = null;
    let activeGenerationId: string | null = null;
    let helloRequestId: string | null = null;
    let negotiated = false;
    let receiveBuffer: Buffer<ArrayBufferLike> = Buffer.alloc(0);
    let nextId = 0;
    let resultDelivered = false;
    let audioLevels: number[] = [];
    let statusRequest: "hello" | "ping" | null = null;
    let statusRequestId: string | null = null;
    let statusBackend: string | null = null;
    let statusModel: string | null = null;
    let statusReceiveBuffer: Buffer<ArrayBufferLike> = Buffer.alloc(0);

    const newId = (kind: string) => `${kind}-${++nextId}`;

    const setStatus = (status: string | undefined) => {
      activeContext?.ui.setStatus(STATUS_KEY, status);
    };

    const setListeningStatus = () => {
      if (!activeContext) return;
      const theme = activeContext.ui.theme;
      const waveform = audioLevels
        .map(
          (level) =>
            WAVEFORM_BARS[
              Math.min(
                WAVEFORM_BARS.length - 1,
                Math.floor(level * WAVEFORM_BARS.length),
              )
            ],
        )
        .join("");
      setStatus(
        `${theme.fg("error", "●")} ${theme.fg("dim", "listening")}  ${theme.fg("accent", waveform)}`,
      );
    };

    const closeSocket = (current: DictateSocket, force: boolean) => {
      try {
        if (force) current.destroy();
        else current.end();
      } catch {
        // Cleanup is best effort and must never deliver text.
      }
    };

    const closeStatusQuery = (forceClose = false) => {
      const current = statusSocket;
      statusSocket = null;
      statusContext = null;
      statusRequest = null;
      statusRequestId = null;
      statusBackend = null;
      statusModel = null;
      statusReceiveBuffer = Buffer.alloc(0);
      if (current) closeSocket(current, forceClose);
    };

    const reset = (forceClose = false) => {
      closeStatusQuery(true);
      const current = socket;
      socket = null;
      activeGenerationId = null;
      helloRequestId = null;
      negotiated = false;
      receiveBuffer = Buffer.alloc(0);
      resultDelivered = false;
      audioLevels = [];
      state = "idle";
      setStatus(undefined);
      activeContext = null;
      if (current) closeSocket(current, forceClose);
    };

    const notifyError = (message: string) => {
      activeContext?.ui.notify(
        boundedText(message, MAX_ERROR_DETAIL_BYTES),
        "error",
      );
    };

    const fail = (message: string) => {
      notifyError(message);
      reset(true);
    };

    const failStatus = (message: string) => {
      const context = statusContext;
      closeStatusQuery(true);
      context?.ui.notify(
        boundedText(message, MAX_ERROR_DETAIL_BYTES),
        "error",
      );
    };

    const sendStatus = (
      current: DictateSocket,
      frame: Record<string, string | number>,
    ): boolean => {
      if (statusSocket !== current) return false;
      const encoded = `${JSON.stringify(frame)}\n`;
      if (utf8Bytes(encoded) - 1 > MAX_FRAME_BYTES) {
        failStatus("Dictation protocol error: outgoing frame is too large.");
        return false;
      }
      try {
        current.write(encoded);
        return true;
      } catch (error) {
        const detail = error instanceof Error ? error.message : "write failed";
        failStatus(`Dictation status request failed: ${detail}`);
        return false;
      }
    };

    const send = (frame: Record<string, string | number>): boolean => {
      if (!socket) return false;
      const encoded = `${JSON.stringify(frame)}\n`;
      if (utf8Bytes(encoded) - 1 > MAX_FRAME_BYTES) {
        fail("Dictation protocol error: outgoing frame is too large.");
        return false;
      }
      try {
        socket.write(encoded);
        return true;
      } catch (error) {
        const detail = error instanceof Error ? error.message : "write failed";
        fail(`Dictation service write failed: ${detail}`);
        return false;
      }
    };

    const sendGenerationCommand = (
      type: "start" | "stop" | "cancel",
    ): boolean => {
      if (!activeGenerationId || !negotiated) return false;
      return send({
        version: PROTOCOL_VERSION,
        type,
        requestId: newId(type),
        generationId: activeGenerationId,
      });
    };

    const protocolError = (detail: string) => {
      fail(`Dictation protocol error: ${detail}`);
    };

    const acceptServiceState = (serviceState: ServiceState): boolean => {
      if (serviceState === "idle") {
        reset();
        return true;
      }
      if (serviceState === "loading" && state === "connecting") {
        setStatus("preparing transcription...");
        return true;
      }
      if (serviceState === "listening" && state === "connecting") {
        state = "recording";
        audioLevels = Array<number>(WAVEFORM_SAMPLES).fill(0);
        setListeningStatus();
        return true;
      }
      if (serviceState === "stopping" && state === "stopping") {
        setStatus("stopping...");
        return true;
      }
      if (serviceState === "transcribing" && state === "stopping") {
        state = "transcribing";
        setStatus("transcribing...");
        return true;
      }
      if (
        (serviceState === "loading" && state === "connecting") ||
        (serviceState === "listening" && state === "recording") ||
        (serviceState === "stopping" && state === "stopping") ||
        (serviceState === "transcribing" && state === "transcribing")
      ) {
        return true;
      }
      return false;
    };

    const handleFrame = (value: unknown, current: DictateSocket) => {
      if (socket !== current) return;
      if (!isRecord(value)) {
        protocolError("frame must be an object.");
        return;
      }
      if (value.version !== PROTOCOL_VERSION) {
        fail("Dictation service protocol version mismatch.");
        return;
      }
      if (typeof value.type !== "string" || !isBoundedId(value.requestId)) {
        protocolError("frame type or request ID is invalid.");
        return;
      }

      if (value.type === "pong") {
        for (const field of ["backend", "model"] as const) {
          if (field in value && !isBoundedDiagnostic(value[field])) {
            protocolError(`diagnostic ${field} is invalid.`);
            return;
          }
        }
        if (
          negotiated ||
          state !== "connecting" ||
          value.requestId !== helloRequestId
        ) {
          protocolError("unexpected hello response.");
          return;
        }
        negotiated = true;
        sendGenerationCommand("start");
        return;
      }
      if (!negotiated) {
        protocolError("generation event received before hello negotiation.");
        return;
      }

      const generationId = value.generationId;
      if (!isBoundedId(generationId)) {
        protocolError("generation ID is invalid.");
        return;
      }
      if (generationId !== activeGenerationId) return;

      if (value.type === "state") {
        const serviceState = value.state;
        if (
          typeof serviceState !== "string" ||
          ![
            "loading",
            "listening",
            "stopping",
            "transcribing",
            "idle",
          ].includes(serviceState)
        ) {
          protocolError("service state is invalid.");
          return;
        }
        if (!acceptServiceState(serviceState as ServiceState)) {
          protocolError("service state transition is invalid.");
        }
        return;
      }

      if (value.type === "level") {
        if (
          state !== "recording" ||
          typeof value.level !== "number" ||
          !Number.isFinite(value.level) ||
          value.level < 0 ||
          value.level > 1
        ) {
          protocolError("audio level is invalid for the current state.");
          return;
        }
        audioLevels = [...audioLevels.slice(1), value.level];
        setListeningStatus();
        return;
      }

      if (value.type === "result") {
        if (state !== "transcribing" || typeof value.transcript !== "string") {
          protocolError("result is invalid for the current state.");
          return;
        }
        if (utf8Bytes(value.transcript) > MAX_TRANSCRIPT_BYTES) {
          protocolError("transcript is too large.");
          return;
        }
        if (resultDelivered) return;
        resultDelivered = true;
        const context = activeContext;
        const transcript = value.transcript;
        reset();
        if (transcript.trim().length === 0) {
          context?.ui.notify("No speech recognized.", "info");
        } else {
          context?.ui.pasteToEditor(transcript);
        }
        return;
      }

      if (value.type === "cancelled") {
        reset();
        return;
      }

      if (value.type === "error") {
        const code = typeof value.code === "string" ? value.code : "service_error";
        const detail = typeof value.detail === "string" ? value.detail : "Dictation failed";
        fail(`Dictation service error (${boundedText(code, 64)}): ${boundedText(detail, 384)}`);
        return;
      }

      protocolError("unknown frame type.");
    };

    const handleStatusFrame = (value: unknown, current: DictateSocket) => {
      if (statusSocket !== current) return;
      if (!isRecord(value)) {
        failStatus("Dictation status protocol error: frame must be an object.");
        return;
      }
      if (value.version !== PROTOCOL_VERSION) {
        failStatus("Dictation status protocol version mismatch.");
        return;
      }
      if (typeof value.type !== "string" || !isBoundedId(value.requestId)) {
        failStatus("Dictation status protocol error: frame is invalid.");
        return;
      }

      if (value.type === "error") {
        const code = typeof value.code === "string" ? value.code : "service_error";
        const detail = typeof value.detail === "string" ? value.detail : "Dictation status failed";
        failStatus(
          `Dictation service error (${boundedText(code, 64)}): ${boundedText(detail, 384)}`,
        );
        return;
      }
      if (value.type !== "pong" || !statusRequestId || value.requestId !== statusRequestId) {
        failStatus("Dictation status protocol error: unexpected response.");
        return;
      }
      for (const field of ["backend", "model"] as const) {
        if (field in value && !isBoundedDiagnostic(value[field])) {
          failStatus(`Dictation status protocol error: diagnostic ${field} is invalid.`);
          return;
        }
      }
      if (isBoundedDiagnostic(value.backend)) statusBackend = value.backend;
      if (isBoundedDiagnostic(value.model)) statusModel = value.model;

      if (statusRequest === "hello") {
        statusRequest = "ping";
        statusRequestId = newId("status-ping");
        sendStatus(current, {
          version: PROTOCOL_VERSION,
          type: "ping",
          requestId: statusRequestId,
        });
        return;
      }
      if (statusRequest !== "ping") {
        failStatus("Dictation status protocol error: status request is invalid.");
        return;
      }

      const backend = statusBackend ?? "unknown";
      const model = statusModel ?? "unknown";
      const context = statusContext;
      closeStatusQuery();
      context?.ui.notify(
        `Dictation backend: ${backend}; model: ${model}`,
        "info",
      );
    };

    const consumeFramedData = (
      chunk: Uint8Array,
      current: DictateSocket,
      isCurrent: () => boolean,
      getBuffer: () => Buffer<ArrayBufferLike>,
      setBuffer: (value: Buffer<ArrayBufferLike>) => void,
      onFrame: (value: unknown, socket: DictateSocket) => void,
      onProtocolError: (detail: string) => void,
    ) => {
      if (!isCurrent()) return;
      let buffer = Buffer.concat([getBuffer(), Buffer.from(chunk)]);
      while (isCurrent()) {
        const newline = buffer.indexOf(0x0a);
        if (newline < 0) {
          setBuffer(buffer);
          if (buffer.length > MAX_FRAME_BYTES) {
            onProtocolError("frame is too large.");
          }
          return;
        }
        if (newline > MAX_FRAME_BYTES) {
          setBuffer(buffer);
          onProtocolError("frame is too large.");
          return;
        }
        const line = buffer.subarray(0, newline);
        buffer = buffer.subarray(newline + 1);
        setBuffer(buffer);
        if (line.length === 0) {
          onProtocolError("empty frame.");
          return;
        }
        let decoded: string;
        try {
          decoded = UTF8_DECODER.decode(line);
        } catch {
          onProtocolError("frame contains invalid UTF-8.");
          return;
        }
        try {
          onFrame(JSON.parse(decoded), current);
        } catch {
          onProtocolError("malformed JSON frame.");
          return;
        }
      }
    };

    const consumeStatusData = (chunk: Uint8Array, current: DictateSocket) =>
      consumeFramedData(
        chunk,
        current,
        () => statusSocket === current,
        () => statusReceiveBuffer,
        (value) => {
          statusReceiveBuffer = value;
        },
        handleStatusFrame,
        (detail) => failStatus(`Dictation status protocol error: ${detail}`),
      );

    const consumeData = (chunk: Uint8Array, current: DictateSocket) =>
      consumeFramedData(
        chunk,
        current,
        () => socket === current,
        () => receiveBuffer,
        (value) => {
          receiveBuffer = value;
        },
        handleFrame,
        protocolError,
      );

    const statusDictation = (ctx: ExtensionContext) => {
      if (state !== "idle") {
        ctx.ui.notify("Dictation is active; status is unavailable.", "info");
        return;
      }
      if (statusSocket) return;

      statusContext = ctx;
      statusRequest = "hello";
      statusRequestId = newId("status-hello");
      statusReceiveBuffer = Buffer.alloc(0);

      let current: DictateSocket;
      try {
        current = socketFactory(socketPath);
      } catch (error) {
        const detail = error instanceof Error ? error.message : "connection failed";
        failStatus(`Unable to query dictation service: ${detail}`);
        return;
      }
      statusSocket = current;
      current.onConnect(() => {
        if (statusSocket !== current || !statusRequestId) return;
        sendStatus(current, {
          version: PROTOCOL_VERSION,
          type: "hello",
          requestId: statusRequestId,
        });
      });
      current.onData((chunk) => consumeStatusData(chunk, current));
      current.onError((error) => {
        if (statusSocket !== current) return;
        failStatus(`Dictation status connection failed: ${error.message}`);
      });
      current.onClose(() => {
        if (statusSocket !== current) return;
        failStatus("Dictation status service disconnected.");
      });
    };

    const startDictation = (ctx: ExtensionContext) => {
      if (state !== "idle") return;
      closeStatusQuery(true);
      state = "connecting";
      activeContext = ctx;
      activeGenerationId = newId("generation");
      helloRequestId = newId("hello");
      negotiated = false;
      resultDelivered = false;
      audioLevels = [];
      receiveBuffer = Buffer.alloc(0);

      let current: DictateSocket;
      try {
        current = socketFactory(socketPath);
      } catch (error) {
        const detail = error instanceof Error ? error.message : "connection failed";
        fail(`Unable to connect to dictation service: ${detail}`);
        return;
      }
      socket = current;
      current.onConnect(() => {
        if (socket !== current || !helloRequestId) return;
        send({
          version: PROTOCOL_VERSION,
          type: "hello",
          requestId: helloRequestId,
        });
      });
      current.onData((chunk) => consumeData(chunk, current));
      current.onError((error) => {
        if (socket !== current) return;
        fail(`Dictation service connection failed: ${error.message}`);
      });
      current.onClose(() => {
        if (socket !== current) return;
        fail("Dictation service disconnected.");
      });
    };

    const stopDictation = () => {
      if (state !== "recording") return;
      state = "stopping";
      sendGenerationCommand("stop");
    };

    const cancelDictation = () => {
      if (state === "idle") return;
      if (negotiated) sendGenerationCommand("cancel");
      reset();
    };

    const toggleDictation = (ctx: ExtensionContext) => {
      if (state === "idle") startDictation(ctx);
      else if (state === "recording") stopDictation();
    };

    pi.registerShortcut("alt+m", {
      description: "Start or stop voice dictation",
      handler: async (ctx) => toggleDictation(ctx),
    });
    pi.registerShortcut("alt+n", {
      description: "Cancel voice dictation",
      handler: async () => cancelDictation(),
    });
    pi.registerCommand("dictate-cancel", {
      description: "Cancel voice dictation",
      handler: async () => cancelDictation(),
    });
    pi.registerCommand("dictate-status", {
      description: "Show the selected dictation backend",
      handler: async (_args, ctx) => statusDictation(ctx),
    });
    pi.on("session_shutdown", async () => {
      closeStatusQuery(true);
      if (state === "idle") return;
      if (negotiated) sendGenerationCommand("cancel");
      reset();
    });
  };
}

export default createDictateExtension();