wraith/frontend/composables/useRdp.ts

import Guacamole from 'guacamole-common-js'
import { useAuthStore } from '~/stores/auth.store'
import { useSessionStore } from '~/stores/session.store'

/**
 * Creates a Guacamole-compatible tunnel that speaks JSON over WebSocket.
 * Does NOT extend Guacamole.Tunnel (its constructor assigns no-op instance
 * properties that shadow subclass methods). Instead, creates a base tunnel
 * and overwrites the methods directly.
 */
function createJsonWsTunnel(wsUrl: string, connectMsg: object) {
  const tunnel = new Guacamole.Tunnel()
  let ws: WebSocket | null = null

  // Custom callbacks for our JSON protocol layer
  let onConnected: ((hostId: number, hostName: string) => void) | null = null
  let onDisconnected: ((reason: string) => void) | null = null
  let onGatewayError: ((message: string) => void) | null = null

  function dispatchInstructions(raw: string) {
    if (!tunnel.oninstruction) return
    let remaining = raw
    while (remaining.length > 0) {
      const semicolonIdx = remaining.indexOf(';')
      if (semicolonIdx === -1) break
      const instruction = remaining.substring(0, semicolonIdx)
      remaining = remaining.substring(semicolonIdx + 1)
      if (!instruction) continue
      const parts: string[] = []
      let pos = 0
      while (pos < instruction.length) {
        const dotIdx = instruction.indexOf('.', pos)
        if (dotIdx === -1) break
        const len = parseInt(instruction.substring(pos, dotIdx), 10)
        if (isNaN(len)) break
        parts.push(instruction.substring(dotIdx + 1, dotIdx + 1 + len))
        pos = dotIdx + 1 + len + 1
      }
      if (parts.length > 0) {
        tunnel.oninstruction(parts[0], parts.slice(1))
      }
    }
  }

  // Override the no-op instance methods
  tunnel.connect = (_data?: string) => {
    console.log('[RDP] Tunnel opening WebSocket:', wsUrl)
    ws = new WebSocket(wsUrl)

    ws.onopen = () => {
      console.log('[RDP] WebSocket open, sending connect handshake')
      ws!.send(JSON.stringify(connectMsg))
    }

    ws.onmessage = (event: MessageEvent) => {
      const msg = JSON.parse(event.data as string)
      switch (msg.type) {
        case 'connected':
          onConnected?.(msg.hostId, msg.hostName)
          break
        case 'guac':
          dispatchInstructions(msg.instruction)
          break
        case 'error':
          onGatewayError?.(msg.message)
          tunnel.onerror?.(new Guacamole.Status(Guacamole.Status.Code.SERVER_ERROR, msg.message))
          break
        case 'disconnected':
          onDisconnected?.(msg.reason)
          break
      }
    }

    ws.onclose = (event: CloseEvent) => {
      console.log('[RDP] WebSocket closed, code:', event.code)
      if (event.code === 4001) {
        tunnel.onerror?.(new Guacamole.Status(Guacamole.Status.Code.CLIENT_FORBIDDEN, 'Unauthorized'))
      }
      onDisconnected?.('WebSocket closed')
    }

    ws.onerror = () => {
      console.error('[RDP] WebSocket error')
      tunnel.onerror?.(new Guacamole.Status(Guacamole.Status.Code.SERVER_ERROR, 'WebSocket error'))
    }
  }

  tunnel.disconnect = () => {
    if (ws && ws.readyState === WebSocket.OPEN) {
      ws.close()
    }
    ws = null
  }

  tunnel.sendMessage = (...elements: any[]) => {
    if (ws?.readyState === WebSocket.OPEN) {
      const instruction = elements.map((e: any) => {
        const s = String(e)
        return `${s.length}.${s}`
      }).join(',') + ';'
      ws.send(JSON.stringify({ type: 'guac', instruction }))
    }
  }

  return {
    tunnel,
    get onConnected() { return onConnected },
    set onConnected(fn) { onConnected = fn },
    get onDisconnected() { return onDisconnected },
    set onDisconnected(fn) { onDisconnected = fn },
    get onGatewayError() { return onGatewayError },
    set onGatewayError(fn) { onGatewayError = fn },
  }
}

// ─── Composable ────────────────────────────────────────────────────────────────

export function useRdp() {
  const auth = useAuthStore()
  const sessions = useSessionStore()

  function connectRdp(
    container: HTMLElement,
    hostId: number,
    hostName: string,
    color: string | null,
    pendingSessionId: string,
    options?: { width?: number; height?: number },
  ) {
    const proto = location.protocol === 'https:' ? 'wss' : 'ws'
    const wsUrl = `${proto}://${location.host}/api/ws/rdp?token=${auth.token}`

    const width = options?.width || container.clientWidth || 1920
    const height = options?.height || container.clientHeight || 1080

    console.log(`[RDP] Connecting to ${wsUrl} for hostId=${hostId} (${width}x${height})`)

    const wrapper = createJsonWsTunnel(wsUrl, {
      type: 'connect',
      hostId,
      width,
      height,
    })

    const client = new Guacamole.Client(wrapper.tunnel)
    const sessionId = pendingSessionId

    wrapper.onConnected = (_resolvedHostId: number, resolvedHostName: string) => {
      console.log(`[RDP] Connected to ${resolvedHostName}`)
      sessions.replaceSession(sessionId, {
        key: sessionId,
        id: sessionId,
        hostId,
        hostName: resolvedHostName || hostName,
        protocol: 'rdp',
        color,
        active: true,
      })
    }

    wrapper.onDisconnected = (reason: string) => {
      console.log(`[RDP] Disconnected: ${reason}`)
      sessions.removeSession(sessionId)
      client.disconnect()
    }

    wrapper.onGatewayError = (message: string) => {
      console.error('[RDP] Gateway error:', message)
    }

    // Handle Guacamole-level errors (NLA failure, auth failure, etc.)
    client.onerror = (status: Guacamole.Status) => {
      const code = status.code
      const msg = status.message || 'Unknown error'
      console.error(`[RDP] Guacamole error: code=${code} message=${msg}`)
      // Surface error via gateway error callback so UI can display it
      wrapper.onGatewayError?.(`RDP connection failed: ${msg}`)
    }

    // Track connection state transitions for diagnostics
    client.onstatechange = (state: number) => {
      const states: Record<number, string> = {
        0: 'IDLE',
        1: 'CONNECTING',
        2: 'WAITING',
        3: 'CONNECTED',
        4: 'DISCONNECTING',
        5: 'DISCONNECTED',
      }
      console.log(`[RDP] State: ${states[state] || state}`)
    }

    // Attach Guacamole display element to container
    const display = client.getDisplay()
    const displayEl = display.getElement()
    container.appendChild(displayEl)

    // Auto-scale the Guacamole display to fit the container
    function fitDisplay() {
      const dw = display.getWidth()
      const dh = display.getHeight()
      if (!dw || !dh) return
      const scale = Math.min(container.clientWidth / dw, container.clientHeight / dh)
      display.scale(scale)
    }

    // Re-fit when guacd sends a sync (display dimensions may have changed)
    const origOnSync = display.onresize
    display.onresize = (w: number, h: number) => {
      origOnSync?.call(display, w, h)
      fitDisplay()
    }

    // Re-fit when the browser container resizes
    const resizeObserver = new ResizeObserver(() => fitDisplay())
    resizeObserver.observe(container)

    // Mouse input — bind to the display element
    const mouse = new Guacamole.Mouse(displayEl)
    mouse.onEach(
      ['mousedown', 'mousemove', 'mouseup'],
      (e: Guacamole.Mouse.Event) => {
        // Scale mouse coordinates back to remote display space
        const scale = display.getScale()
        const scaledState = new Guacamole.Mouse.State(
          e.state.x / scale,
          e.state.y / scale,
          e.state.left,
          e.state.middle,
          e.state.right,
          e.state.up,
          e.state.down,
        )
        client.sendMouseState(scaledState)
      },
    )

    // Keyboard input
    const keyboard = new Guacamole.Keyboard(document)
    keyboard.onkeydown = (keysym: number) => client.sendKeyEvent(1, keysym)
    keyboard.onkeyup = (keysym: number) => client.sendKeyEvent(0, keysym)

    // Initiate the WebSocket connection
    client.connect()

    function disconnect() {
      resizeObserver.disconnect()
      keyboard.onkeydown = null
      keyboard.onkeyup = null
      client.disconnect()
      sessions.removeSession(sessionId)
    }

    function sendClipboardText(text: string) {
      const stream = client.createClipboardStream('text/plain')
      const writer = new Guacamole.StringWriter(stream)
      writer.sendText(text)
      writer.sendEnd()
    }

    return { client, sessionId, disconnect, sendClipboardText }
  }

  return { connectRdp }
}