remotelink-docker/app/viewer/[sessionId]/page.tsx

'use client'

import { useEffect, useState, useRef, useCallback } from 'react'
import { useParams, useRouter, useSearchParams } from 'next/navigation'
import { Button } from '@/components/ui/button'
import { Monitor, Loader2, AlertCircle, WifiOff } from 'lucide-react'
import { ViewerToolbar } from '@/components/viewer/toolbar'
import { ConnectionStatus } from '@/components/viewer/connection-status'
import { FileTransferPanel } from '@/components/viewer/file-transfer-panel'
import { ChatPanel } from '@/components/viewer/chat-panel'

interface Session {
  id: string
  machineId: string | null
  machineName: string | null
  startedAt: string
  endedAt: string | null
  connectionType: string | null
}

interface MonitorInfo {
  index: number
  width: number
  height: number
  primary: boolean
}

type ConnectionState = 'connecting' | 'waiting' | 'connected' | 'disconnected' | 'error'

export default function ViewerPage() {
  const params = useParams()
  const searchParams = useSearchParams()
  const router = useRouter()
  const sessionId = params.sessionId as string
  const viewerToken = searchParams.get('token')

  const [session, setSession] = useState<Session | null>(null)
  const [connectionState, setConnectionState] = useState<ConnectionState>('connecting')
  const [isFullscreen, setIsFullscreen] = useState(false)
  const [showToolbar, setShowToolbar] = useState(true)
  const [quality, setQuality] = useState<'high' | 'medium' | 'low'>('high')
  const [error, setError] = useState<string | null>(null)
  const [statusMsg, setStatusMsg] = useState('Connecting to relay…')
  const [fps, setFps] = useState(0)
  const [monitors, setMonitors] = useState<MonitorInfo[]>([])
  const [activeMonitor, setActiveMonitor] = useState(0)
  const [showFileTransfer, setShowFileTransfer] = useState(false)
  const [showChat, setShowChat] = useState(false)
  const [lastAgentMsg, setLastAgentMsg] = useState<Record<string, unknown> | null>(null)

  const containerRef = useRef<HTMLDivElement>(null)
  const canvasRef = useRef<HTMLCanvasElement>(null)
  const wsRef = useRef<WebSocket | null>(null)
  const fpsCounterRef = useRef({ frames: 0, lastTime: Date.now() })
  const reconnectTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null)
  const connectionStateRef = useRef<ConnectionState>('connecting')

  // Keep ref in sync so closures can read current value
  useEffect(() => { connectionStateRef.current = connectionState }, [connectionState])

  // ── Load session info ──────────────────────────────────────────────────────
  useEffect(() => {
    fetch(`/api/sessions/${sessionId}`)
      .then((r) => r.json())
      .then((data) => {
        if (!data.session) { setError('Session not found'); setConnectionState('error'); return }
        if (data.session.endedAt) { setError('This session has ended'); setConnectionState('disconnected'); return }
        setSession(data.session)
      })
      .catch(() => { setError('Failed to load session'); setConnectionState('error') })
  }, [sessionId])

  // ── WebSocket connection ───────────────────────────────────────────────────
  const connectWS = useCallback(() => {
    if (!viewerToken) return
    if (reconnectTimerRef.current) { clearTimeout(reconnectTimerRef.current); reconnectTimerRef.current = null }

    const relayHost = process.env.NEXT_PUBLIC_RELAY_URL || `${window.location.hostname}:8765`
    const proto = window.location.protocol === 'https:' ? 'wss' : 'ws'
    const wsUrl = `${proto}://${relayHost}/ws/viewer?session_id=${sessionId}&viewer_token=${viewerToken}`

    setConnectionState('connecting')
    setStatusMsg('Connecting to relay…')

    const ws = new WebSocket(wsUrl)
    ws.binaryType = 'arraybuffer'
    wsRef.current = ws

    ws.onopen = () => { setStatusMsg('Waiting for agent…'); setConnectionState('waiting') }

    ws.onmessage = (evt) => {
      if (evt.data instanceof ArrayBuffer) {
        renderFrame(evt.data)
      } else {
        try {
          const msg = JSON.parse(evt.data as string)
          handleRelayMessage(msg)
        } catch { /* ignore */ }
      }
    }

    ws.onclose = (evt) => {
      wsRef.current = null
      if (connectionStateRef.current !== 'disconnected' && evt.code !== 4001) {
        setConnectionState('waiting')
        setStatusMsg('Connection lost — reconnecting…')
        reconnectTimerRef.current = setTimeout(connectWS, 3000)
      }
    }

    ws.onerror = () => { setStatusMsg('Relay connection failed') }
  }, [sessionId, viewerToken]) // eslint-disable-line react-hooks/exhaustive-deps

  useEffect(() => {
    if (session && viewerToken) connectWS()
    return () => {
      wsRef.current?.close()
      if (reconnectTimerRef.current) clearTimeout(reconnectTimerRef.current)
    }
  }, [session, viewerToken, connectWS])

  const handleRelayMessage = (msg: Record<string, unknown>) => {
    switch (msg.type) {
      case 'agent_connected':
        setConnectionState('connected')
        setStatusMsg('')
        break
      case 'agent_offline':
        setConnectionState('waiting')
        setStatusMsg('Agent is offline — waiting for connection…')
        break
      case 'agent_disconnected':
        setConnectionState('waiting')
        setStatusMsg('Agent disconnected — waiting for reconnect…')
        break
      case 'monitor_list':
        setMonitors((msg.monitors as MonitorInfo[]) ?? [])
        break
      case 'clipboard_content':
        // Write agent clipboard to browser clipboard
        navigator.clipboard.writeText(msg.content as string).catch(() => {})
        break
      case 'file_chunk':
      case 'file_list':
      case 'chat_message':
        setLastAgentMsg(msg)
        break
    }
  }

  // ── Frame rendering ────────────────────────────────────────────────────────
  const renderFrame = useCallback((data: ArrayBuffer) => {
    const canvas = canvasRef.current
    if (!canvas) return
    const blob = new Blob([data], { type: 'image/jpeg' })
    const url = URL.createObjectURL(blob)
    const img = new Image()
    img.onload = () => {
      const ctx = canvas.getContext('2d')
      if (!ctx) return
      if (canvas.width !== img.width || canvas.height !== img.height) {
        canvas.width = img.width
        canvas.height = img.height
      }
      ctx.drawImage(img, 0, 0)
      URL.revokeObjectURL(url)
      const counter = fpsCounterRef.current
      counter.frames++
      const now = Date.now()
      if (now - counter.lastTime >= 1000) {
        setFps(counter.frames)
        counter.frames = 0
        counter.lastTime = now
      }
    }
    img.src = url
  }, [])

  // ── Input forwarding ───────────────────────────────────────────────────────
  const sendEvent = useCallback((event: Record<string, unknown>) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(event))
    }
  }, [])

  const getCanvasCoords = useCallback((e: React.MouseEvent<HTMLCanvasElement>) => {
    const canvas = canvasRef.current!
    const rect = canvas.getBoundingClientRect()
    return {
      x: Math.round(((e.clientX - rect.left) * canvas.width) / rect.width),
      y: Math.round(((e.clientY - rect.top) * canvas.height) / rect.height),
    }
  }, [])

  const handleMouseMove = useCallback((e: React.MouseEvent<HTMLCanvasElement>) => {
    if (connectionState !== 'connected') return
    const { x, y } = getCanvasCoords(e)
    sendEvent({ type: 'mouse_move', x, y })
  }, [connectionState, sendEvent, getCanvasCoords])

  const handleMouseClick = useCallback((e: React.MouseEvent<HTMLCanvasElement>) => {
    if (connectionState !== 'connected') return
    const { x, y } = getCanvasCoords(e)
    sendEvent({ type: 'mouse_click', button: e.button === 2 ? 'right' : 'left', double: e.detail === 2, x, y })
  }, [connectionState, sendEvent, getCanvasCoords])

  const handleMouseScroll = useCallback((e: React.WheelEvent<HTMLCanvasElement>) => {
    if (connectionState !== 'connected') return
    e.preventDefault()
    sendEvent({ type: 'mouse_scroll', dx: e.deltaX, dy: -e.deltaY })
  }, [connectionState, sendEvent])

  const handleKeyDown = useCallback((e: KeyboardEvent) => {
    if (connectionState !== 'connected') return
    e.preventDefault()
    const specialKeys = [
      'Enter', 'Escape', 'Tab', 'Backspace', 'Delete', 'Home', 'End',
      'PageUp', 'PageDown', 'ArrowUp', 'ArrowDown', 'ArrowLeft', 'ArrowRight',
      'F1','F2','F3','F4','F5','F6','F7','F8','F9','F10','F11','F12',
    ]
    if (specialKeys.includes(e.key)) {
      sendEvent({ type: 'key_special', key: e.key.toLowerCase().replace('arrow', '') })
    } else if (e.key.length === 1) {
      sendEvent({ type: 'key_press', key: e.key })
    }
  }, [connectionState, sendEvent])

  useEffect(() => {
    document.addEventListener('keydown', handleKeyDown)
    return () => document.removeEventListener('keydown', handleKeyDown)
  }, [handleKeyDown])

  // ── Toolbar actions ────────────────────────────────────────────────────────
  const handleCtrlAltDel = useCallback(() => {
    sendEvent({ type: 'exec_key_combo', keys: ['ctrl_l', 'alt_l', 'delete'] })
  }, [sendEvent])

  const handleClipboardPaste = useCallback(async () => {
    try {
      const text = await navigator.clipboard.readText()
      sendEvent({ type: 'clipboard_paste', content: text })
    } catch {
      // clipboard permission denied — ignore
    }
  }, [sendEvent])

  const handleMonitorChange = useCallback((index: number) => {
    setActiveMonitor(index)
    sendEvent({ type: 'set_monitor', index })
  }, [sendEvent])

  const handleQualityChange = useCallback((q: 'high' | 'medium' | 'low') => {
    setQuality(q)
    sendEvent({ type: 'set_quality', quality: q })
  }, [sendEvent])

  // ── Fullscreen / toolbar auto-hide ─────────────────────────────────────────
  const toggleFullscreen = async () => {
    if (!containerRef.current) return
    if (!document.fullscreenElement) {
      await containerRef.current.requestFullscreen()
      setIsFullscreen(true)
    } else {
      await document.exitFullscreen()
      setIsFullscreen(false)
    }
  }

  useEffect(() => {
    if (isFullscreen && connectionState === 'connected') {
      const timer = setTimeout(() => setShowToolbar(false), 3000)
      return () => clearTimeout(timer)
    } else {
      setShowToolbar(true)
    }
  }, [isFullscreen, connectionState])

  // ── End session ────────────────────────────────────────────────────────────
  const endSession = async () => {
    wsRef.current?.close()
    if (session) {
      const duration = Math.floor((Date.now() - new Date(session.startedAt).getTime()) / 1000)
      await fetch(`/api/sessions/${session.id}`, {
        method: 'PATCH',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ endedAt: new Date().toISOString(), durationSeconds: duration }),
      })
    }
    setConnectionState('disconnected')
    router.push('/dashboard/sessions')
  }

  // ── Render ─────────────────────────────────────────────────────────────────
  if (error) {
    return (
      <div className="min-h-svh bg-background flex items-center justify-center">
        <div className="text-center space-y-4">
          <AlertCircle className="h-12 w-12 text-destructive mx-auto" />
          <h1 className="text-xl font-semibold">{error}</h1>
          <Button onClick={() => router.push('/dashboard')}>Return to Dashboard</Button>
        </div>
      </div>
    )
  }

  return (
    <div
      ref={containerRef}
      className="min-h-svh bg-black flex flex-col"
      onMouseMove={() => isFullscreen && setShowToolbar(true)}
    >
      <div className={`transition-opacity duration-300 ${showToolbar ? 'opacity-100' : 'opacity-0 pointer-events-none'}`}>
        <ViewerToolbar
          session={session}
          connectionState={connectionState}
          isFullscreen={isFullscreen}
          quality={quality}
          isMuted={true}
          monitors={monitors}
          activeMonitor={activeMonitor}
          onToggleFullscreen={toggleFullscreen}
          onQualityChange={handleQualityChange}
          onToggleMute={() => {}}
          onDisconnect={endSession}
          onReconnect={connectWS}
          onCtrlAltDel={handleCtrlAltDel}
          onClipboardPaste={handleClipboardPaste}
          onMonitorChange={handleMonitorChange}
          onToggleFileTransfer={() => setShowFileTransfer(v => !v)}
          onToggleChat={() => setShowChat(v => !v)}
          onWakeOnLan={session?.machineId ? async () => {
            await fetch('/api/machines/wol', {
              method: 'POST',
              headers: { 'Content-Type': 'application/json' },
              body: JSON.stringify({ machineId: session.machineId }),
            })
          } : undefined}
        />
      </div>

      <div className="flex flex-1 overflow-hidden">
        {/* Main canvas area */}
        <div className="flex-1 flex items-center justify-center relative">
          {(connectionState === 'connecting' || connectionState === 'waiting') && (
            <div className="absolute inset-0 flex items-center justify-center bg-background/80 z-10">
              <div className="text-center space-y-4">
                {connectionState === 'waiting'
                  ? <WifiOff className="h-12 w-12 text-muted-foreground mx-auto" />
                  : <Loader2 className="h-12 w-12 animate-spin text-primary mx-auto" />}
                <div>
                  <p className="font-semibold">{statusMsg}</p>
                  <p className="text-sm text-muted-foreground">{session?.machineName}</p>
                </div>
              </div>
            </div>
          )}

          <canvas
            ref={canvasRef}
            className="max-w-full max-h-full object-contain cursor-crosshair"
            style={{
              display: connectionState === 'connected' ? 'block' : 'none',
              imageRendering: quality === 'low' ? 'pixelated' : 'auto',
            }}
            onMouseMove={handleMouseMove}
            onClick={handleMouseClick}
            onWheel={handleMouseScroll}
            onContextMenu={(e) => { e.preventDefault(); handleMouseClick(e) }}
          />

          {connectionState === 'disconnected' && (
            <div className="text-center space-y-4">
              <Monitor className="h-16 w-16 text-muted-foreground mx-auto" />
              <div>
                <p className="font-semibold">Session Ended</p>
                <p className="text-sm text-muted-foreground">The remote connection has been closed</p>
              </div>
              <Button onClick={() => router.push('/dashboard')}>Return to Dashboard</Button>
            </div>
          )}
        </div>

        {/* Sidebars */}
        {showFileTransfer && connectionState === 'connected' && (
          <FileTransferPanel
            onClose={() => setShowFileTransfer(false)}
            sendEvent={sendEvent}
            incomingMessage={lastAgentMsg}
          />
        )}
        {showChat && connectionState === 'connected' && (
          <ChatPanel
            onClose={() => setShowChat(false)}
            sendEvent={sendEvent}
            incomingMessage={lastAgentMsg}
          />
        )}
      </div>

      {connectionState === 'connected' && <ConnectionStatus quality={quality} fps={fps} />}
    </div>
  )
}