wraith/src/composables/useRdp.ts

import { ref, onBeforeUnmount } from "vue";
import type { Ref } from "vue";
import { invoke } from "@tauri-apps/api/core";

/**
 * RDP mouse event flags — match the Rust constants in src-tauri/src/rdp/input.rs
 */
export const MouseFlag = {
  Move: 0x0800,
  Button1: 0x1000, // Left
  Button2: 0x2000, // Right
  Button3: 0x4000, // Middle
  Down: 0x8000,
  Wheel: 0x0200,
  WheelNeg: 0x0100,
  HWheel: 0x0400,
} as const;

/**
 * JavaScript KeyboardEvent.code → RDP scancode mapping.
 * Mirrors the Rust ScancodeMap in src-tauri/src/rdp/input.rs.
 */
export const ScancodeMap: Record<string, number> = {
  // Row 0: Escape + Function keys
  Escape: 0x0001,
  F1: 0x003b,
  F2: 0x003c,
  F3: 0x003d,
  F4: 0x003e,
  F5: 0x003f,
  F6: 0x0040,
  F7: 0x0041,
  F8: 0x0042,
  F9: 0x0043,
  F10: 0x0044,
  F11: 0x0057,
  F12: 0x0058,

  // Row 1: Number row
  Backquote: 0x0029,
  Digit1: 0x0002,
  Digit2: 0x0003,
  Digit3: 0x0004,
  Digit4: 0x0005,
  Digit5: 0x0006,
  Digit6: 0x0007,
  Digit7: 0x0008,
  Digit8: 0x0009,
  Digit9: 0x000a,
  Digit0: 0x000b,
  Minus: 0x000c,
  Equal: 0x000d,
  Backspace: 0x000e,

  // Row 2: QWERTY
  Tab: 0x000f,
  KeyQ: 0x0010,
  KeyW: 0x0011,
  KeyE: 0x0012,
  KeyR: 0x0013,
  KeyT: 0x0014,
  KeyY: 0x0015,
  KeyU: 0x0016,
  KeyI: 0x0017,
  KeyO: 0x0018,
  KeyP: 0x0019,
  BracketLeft: 0x001a,
  BracketRight: 0x001b,
  Backslash: 0x002b,

  // Row 3: Home row
  CapsLock: 0x003a,
  KeyA: 0x001e,
  KeyS: 0x001f,
  KeyD: 0x0020,
  KeyF: 0x0021,
  KeyG: 0x0022,
  KeyH: 0x0023,
  KeyJ: 0x0024,
  KeyK: 0x0025,
  KeyL: 0x0026,
  Semicolon: 0x0027,
  Quote: 0x0028,
  Enter: 0x001c,

  // Row 4: Bottom row
  ShiftLeft: 0x002a,
  KeyZ: 0x002c,
  KeyX: 0x002d,
  KeyC: 0x002e,
  KeyV: 0x002f,
  KeyB: 0x0030,
  KeyN: 0x0031,
  KeyM: 0x0032,
  Comma: 0x0033,
  Period: 0x0034,
  Slash: 0x0035,
  ShiftRight: 0x0036,

  // Row 5: Bottom modifiers + space
  ControlLeft: 0x001d,
  MetaLeft: 0xe05b,
  AltLeft: 0x0038,
  Space: 0x0039,
  AltRight: 0xe038,
  MetaRight: 0xe05c,
  ContextMenu: 0xe05d,
  ControlRight: 0xe01d,

  // Navigation cluster
  PrintScreen: 0xe037,
  ScrollLock: 0x0046,
  Pause: 0x0045,
  Insert: 0xe052,
  Home: 0xe047,
  PageUp: 0xe049,
  Delete: 0xe053,
  End: 0xe04f,
  PageDown: 0xe051,

  // Arrow keys
  ArrowUp: 0xe048,
  ArrowLeft: 0xe04b,
  ArrowDown: 0xe050,
  ArrowRight: 0xe04d,

  // Numpad
  NumLock: 0x0045,
  NumpadDivide: 0xe035,
  NumpadMultiply: 0x0037,
  NumpadSubtract: 0x004a,
  Numpad7: 0x0047,
  Numpad8: 0x0048,
  Numpad9: 0x0049,
  NumpadAdd: 0x004e,
  Numpad4: 0x004b,
  Numpad5: 0x004c,
  Numpad6: 0x004d,
  Numpad1: 0x004f,
  Numpad2: 0x0050,
  Numpad3: 0x0051,
  NumpadEnter: 0xe01c,
  Numpad0: 0x0052,
  NumpadDecimal: 0x0053,
};

/**
 * Look up the RDP scancode for a JS KeyboardEvent.code string.
 */
export function jsKeyToScancode(code: string): number | null {
  return ScancodeMap[code] ?? null;
}

export interface UseRdpReturn {
  /** Whether the RDP session is connected (first frame received) */
  connected: Ref<boolean>;
  /** Whether keyboard capture is enabled */
  keyboardGrabbed: Ref<boolean>;
  /** Whether clipboard sync is enabled */
  clipboardSync: Ref<boolean>;
  /** Fetch and render the dirty region directly to a canvas context */
  fetchAndRender: (sessionId: string, width: number, height: number, ctx: CanvasRenderingContext2D) => Promise<boolean>;
  /** Send a mouse event to the backend */
  sendMouse: (sessionId: string, x: number, y: number, flags: number) => void;
  /** Send a key event to the backend */
  sendKey: (sessionId: string, code: string, pressed: boolean) => void;
  /** Send clipboard text to the remote session */
  sendClipboard: (sessionId: string, text: string) => void;
  /** Start the frame rendering loop targeting ~30fps */
  startFrameLoop: (
    sessionId: string,
    canvas: HTMLCanvasElement,
    width: number,
    height: number,
  ) => void;
  /** Stop the frame rendering loop */
  stopFrameLoop: () => void;
  /** Toggle keyboard grab */
  toggleKeyboardGrab: () => void;
  /** Toggle clipboard sync */
  toggleClipboardSync: () => void;
}

/**
 * Composable that manages an RDP session's rendering and input.
 *
 * Uses Tauri's invoke() to call Rust commands:
 *   rdp_get_frame   → raw RGBA ArrayBuffer (binary IPC)
 *   rdp_send_mouse  → fire-and-forget
 *   rdp_send_key    → fire-and-forget
 *   rdp_send_clipboard → fire-and-forget
 */
export function useRdp(): UseRdpReturn {
  const connected = ref(false);
  const keyboardGrabbed = ref(true);
  const clipboardSync = ref(false);

  let animFrameId: number | null = null;
  let unlistenFrame: (() => void) | null = null;

  /**
   * Fetch the dirty region from the Rust RDP backend and apply it to the canvas.
   *
   * Binary format from backend: 8-byte header + pixel data
   *   Header: [x: u16, y: u16, w: u16, h: u16] (little-endian)
   *   If header is all zeros → full frame (width*height*4 bytes)
   *   If header is non-zero → dirty rectangle (w*h*4 bytes)
   *
   * Returns true if a frame was rendered, false if nothing changed.
   */
  async function fetchAndRender(
    sessionId: string,
    width: number,
    height: number,
    ctx: CanvasRenderingContext2D,
  ): Promise<boolean> {
    let raw: ArrayBuffer;
    try {
      raw = await invoke<ArrayBuffer>("rdp_get_frame", { sessionId });
    } catch {
      return false;
    }

    if (!raw || raw.byteLength <= 8) return false;

    const view = new DataView(raw);
    const rx = view.getUint16(0, true);
    const ry = view.getUint16(2, true);
    const rw = view.getUint16(4, true);
    const rh = view.getUint16(6, true);
    const pixelData = new Uint8ClampedArray(raw, 8);

    if (rx === 0 && ry === 0 && rw === 0 && rh === 0) {
      // Full frame
      const expected = width * height * 4;
      if (pixelData.length !== expected) return false;
      ctx.putImageData(new ImageData(pixelData, width, height), 0, 0);
    } else {
      // Dirty rectangle — apply at offset
      const expected = rw * rh * 4;
      if (pixelData.length !== expected) return false;
      ctx.putImageData(new ImageData(pixelData, rw, rh), rx, ry);
    }

    return true;
  }

  /**
   * Send a mouse event to the remote session.
   * Calls Rust rdp_send_mouse(sessionId, x, y, flags).
   * Fire-and-forget — mouse events are best-effort.
   */
  function sendMouse(
    sessionId: string,
    x: number,
    y: number,
    flags: number,
  ): void {
    invoke("rdp_send_mouse", { sessionId, x, y, flags }).catch(
      (err: unknown) => {
        console.warn("[useRdp] sendMouse failed:", err);
      },
    );
  }

  /**
   * Send a key event, mapping the JS KeyboardEvent.code to an RDP scancode.
   * Calls Rust rdp_send_key(sessionId, scancode, pressed).
   * Unmapped keys are silently dropped — not every JS key has an RDP scancode.
   */
  function sendKey(sessionId: string, code: string, pressed: boolean): void {
    const scancode = jsKeyToScancode(code);
    if (scancode === null) return;

    invoke("rdp_send_key", { sessionId, scancode, pressed }).catch(
      (err: unknown) => {
        console.warn("[useRdp] sendKey failed:", err);
      },
    );
  }

  /**
   * Send clipboard text to the remote RDP session.
   * Calls Rust rdp_send_clipboard(sessionId, text).
   */
  function sendClipboard(sessionId: string, text: string): void {
    invoke("rdp_send_clipboard", { sessionId, text }).catch(
      (err: unknown) => {
        console.warn("[useRdp] sendClipboard failed:", err);
      },
    );
  }

  /**
   * Start the rendering loop. Fetches frames and draws them on the canvas.
   *
   * Targets ~30fps by skipping every other rAF tick (rAF fires at ~60fps).
   * Sets connected = true as soon as the loop starts — the overlay dismisses
   * on first successful frame render.
   */
  function startFrameLoop(
    sessionId: string,
    canvas: HTMLCanvasElement,
    width: number,
    height: number,
  ): void {
    const ctx = canvas.getContext("2d");
    if (!ctx) return;

    canvas.width = width;
    canvas.height = height;

    let fetchPending = false;
    let rafScheduled = false;

    // Fetch and render dirty region when backend signals new frame data.
    // Uses rAF to coalesce rapid events into one fetch per display frame.
    function scheduleFrameFetch(): void {
      if (rafScheduled) return;
      rafScheduled = true;
      animFrameId = requestAnimationFrame(async () => {
        rafScheduled = false;
        if (fetchPending) return;
        fetchPending = true;
        if (!ctx) return;
        const rendered = await fetchAndRender(sessionId, width, height, ctx);
        fetchPending = false;
        if (rendered && !connected.value) connected.value = true;
      });
    }

    // Listen for frame events from the backend (push model)
    import("@tauri-apps/api/event").then(({ listen }) => {
      listen(`rdp:frame:${sessionId}`, () => {
        scheduleFrameFetch();
      }).then((unlisten) => {
        unlistenFrame = unlisten;
      });
    });

    // Initial poll in case frames arrived before listener was set up
    scheduleFrameFetch();
  }

  /**
   * Stop the rendering loop and reset connected state.
   */
  function stopFrameLoop(): void {
    if (animFrameId !== null) {
      cancelAnimationFrame(animFrameId);
      animFrameId = null;
    }
    if (unlistenFrame !== null) {
      unlistenFrame();
      unlistenFrame = null;
    }
    connected.value = false;
  }

  function toggleKeyboardGrab(): void {
    keyboardGrabbed.value = !keyboardGrabbed.value;
  }

  function toggleClipboardSync(): void {
    clipboardSync.value = !clipboardSync.value;
  }

  onBeforeUnmount(() => {
    stopFrameLoop();
  });

  return {
    connected,
    keyboardGrabbed,
    clipboardSync,
    fetchAndRender,
    sendMouse,
    sendKey,
    sendClipboard,
    startFrameLoop,
    stopFrameLoop,
    toggleKeyboardGrab,
    toggleClipboardSync,
  };
}