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 { Maximize2, Minimize2, Monitor, Loader2, AlertCircle, WifiOff } from 'lucide-react'
import { ViewerToolbar } from '@/components/viewer/toolbar'
import { ConnectionStatus } from '@/components/viewer/connection-status'

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

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
  // viewerToken is passed as a query param from the connect flow
  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 [isMuted, setIsMuted] = useState(true)
  const [error, setError] = useState<string | null>(null)
  const [statusMsg, setStatusMsg] = useState('Connecting to relay…')
  const [fps, setFps] = useState(0)

  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)

  // ── 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

    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) {
        // Binary = JPEG frame
        renderFrame(evt.data)
      } else {
        try {
          const msg = JSON.parse(evt.data as string)
          handleRelayMessage(msg)
        } catch { /* ignore */ }
      }
    }

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

    ws.onerror = () => {
      setStatusMsg('Relay connection failed')
    }
  }, [sessionId, viewerToken])

  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
    }
  }

  // ── 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)

      // FPS counter
      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])

  // ── 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)
    }
  }, [isFullscreen, connectionState, showToolbar])

  // ── 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'}`}>
        <ViewerToolbar
          session={session}
          connectionState={connectionState}
          isFullscreen={isFullscreen}
          quality={quality}
          isMuted={isMuted}
          onToggleFullscreen={toggleFullscreen}
          onQualityChange={(q) => {
            setQuality(q)
            sendEvent({ type: 'set_quality', quality: q })
          }}
          onToggleMute={() => setIsMuted(!isMuted)}
          onDisconnect={endSession}
          onReconnect={connectWS}
        />
      </div>

      <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>

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