wraith/src/composables/useRdp.ts

import { ref, onBeforeUnmount } 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: ReturnType<typeof ref<boolean>>;
  /** Whether keyboard capture is enabled */
  keyboardGrabbed: ReturnType<typeof ref<boolean>>;
  /** Whether clipboard sync is enabled */
  clipboardSync: ReturnType<typeof ref<boolean>>;
  /** Fetch the current frame as RGBA ImageData */
  fetchFrame: (sessionId: string, width: number, height: number) => Promise<ImageData | null>;
  /** 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   → base64 RGBA string
 *   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 frameCount = 0;

  /**
   * Fetch the current frame from the Rust RDP backend.
   *
   * rdp_get_frame returns raw RGBA bytes (width*height*4) serialised as a
   * base64 string over Tauri's IPC bridge. We decode it to Uint8ClampedArray
   * and wrap in an ImageData for putImageData().
   */
  async function fetchFrame(
    sessionId: string,
    width: number,
    height: number,
  ): Promise<ImageData | null> {
    let raw: string;
    try {
      raw = await invoke<string>("rdp_get_frame", { sessionId });
    } catch {
      // Session may not be connected yet or backend returned an error — skip frame
      return null;
    }

    if (!raw || raw.length === 0) return null;

    // Decode base64 → binary string → Uint8ClampedArray
    const binaryStr = atob(raw);
    const bytes = new Uint8ClampedArray(binaryStr.length);
    for (let i = 0; i < binaryStr.length; i++) {
      bytes[i] = binaryStr.charCodeAt(i);
    }

    // Validate: RGBA requires exactly width * height * 4 bytes
    const expected = width * height * 4;
    if (bytes.length !== expected) {
      console.warn(
        `[useRdp] Frame size mismatch: got ${bytes.length}, expected ${expected}`,
      );
      return null;
    }

    return new ImageData(bytes, width, height);
  }

  /**
   * 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;

    function renderLoop(): void {
      frameCount++;

      // Throttle to ~30fps by skipping odd-numbered rAF ticks
      if (frameCount % 2 === 0) {
        fetchFrame(sessionId, width, height).then((imageData) => {
          if (imageData && ctx) {
            ctx.putImageData(imageData, 0, 0);
            // Mark connected on first successful frame
            if (!connected.value) {
              connected.value = true;
            }
          }
        });
      }

      animFrameId = requestAnimationFrame(renderLoop);
    }

    animFrameId = requestAnimationFrame(renderLoop);
  }

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

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

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

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

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