OpenSquawk/server/api/flightlab/ws.ts

// server/api/flightlab/ws.ts
import { defineWebSocketHandler } from 'h3'
import { flightlabTelemetryStore } from '../../utils/flightlabTelemetry'

interface FlightLabSession {
  code: string
  scenarioId: string
  currentPhaseId: string
  isPaused: boolean
  startedAt: number
  participantConnected: boolean
  instructorConnected: boolean
  history: Array<{ phaseId: string; buttonId: string; timestamp: number }>
  peers: Map<string, { role: 'instructor' | 'participant'; peer: any; userId?: string }>
  /** User IDs subscribed to telemetry for this session */
  telemetryUserIds: Set<string>
}

const GLOBAL_SESSION_CODE = 'GLOBAL'
const sessions = new Map<string, FlightLabSession>()

// Map userId → session code for telemetry routing
const userSessionMap = new Map<string, string>()

// Subscribe to telemetry store updates and relay to WebSocket clients
flightlabTelemetryStore.subscribe((userId, data) => {
  const sessionCode = userSessionMap.get(userId)
  if (!sessionCode) return
  const session = sessions.get(sessionCode)
  if (!session) return
  broadcastToSession(session, { type: 'telemetry', data })
})

function getOrCreateGlobalSession(scenarioId = 'takeoff-eddf'): FlightLabSession {
  const existing = sessions.get(GLOBAL_SESSION_CODE)
  if (existing) {
    existing.scenarioId = scenarioId
    return existing
  }

  const session: FlightLabSession = {
    code: GLOBAL_SESSION_CODE,
    scenarioId,
    currentPhaseId: 'welcome',
    isPaused: false,
    startedAt: Date.now(),
    participantConnected: false,
    instructorConnected: false,
    history: [],
    peers: new Map(),
    telemetryUserIds: new Set(),
  }
  sessions.set(GLOBAL_SESSION_CODE, session)
  return session
}

function broadcastToSession(session: FlightLabSession, event: any, excludePeerId?: string) {
  for (const [id, p] of session.peers) {
    if (id !== excludePeerId) {
      try { p.peer.send(JSON.stringify(event)) } catch {}
    }
  }
}

function getSessionState(session: FlightLabSession) {
  return {
    sessionCode: session.code,
    scenarioId: session.scenarioId,
    currentPhaseId: session.currentPhaseId,
    isPaused: session.isPaused,
    startedAt: session.startedAt,
    participantConnected: session.participantConnected,
    history: session.history,
  }
}

function findSessionByPeer(peerId: string): FlightLabSession | undefined {
  for (const session of sessions.values()) {
    if (session.peers.has(peerId)) return session
  }
}

export default defineWebSocketHandler({
  open(_peer) {
    // Connection opened, wait for join/create message
  },

  message(peer, msg) {
    let data: any
    try {
      data = JSON.parse(typeof msg === 'string' ? msg : msg.toString())
    } catch {
      peer.send(JSON.stringify({ type: 'error', message: 'Invalid JSON' }))
      return
    }

    const peerId = (peer as any).id ?? String(peer)

    switch (data.type) {
      case 'create-session': {
        const session = getOrCreateGlobalSession(data.scenarioId ?? 'takeoff-eddf')
        session.peers.set(peerId, { role: 'instructor', peer })
        session.instructorConnected = true
        peer.send(JSON.stringify({ type: 'session-created', code: session.code, state: getSessionState(session) }))
        broadcastToSession(session, { type: 'peer-joined', role: 'instructor' }, peerId)
        break
      }

      case 'join-session': {
        const requestedCode = typeof data.code === 'string' ? data.code.toUpperCase() : null
        const session = requestedCode ? sessions.get(requestedCode) : getOrCreateGlobalSession(data.scenarioId ?? 'takeoff-eddf')
        if (!session) {
          peer.send(JSON.stringify({ type: 'error', message: 'Session nicht gefunden' }))
          return
        }
        const role = data.role ?? 'participant'
        session.peers.set(peerId, { role, peer })
        if (role === 'participant') session.participantConnected = true
        peer.send(JSON.stringify({ type: 'session-joined', role, state: getSessionState(session) }))
        broadcastToSession(session, { type: 'peer-joined', role }, peerId)
        break
      }

      case 'participant-action': {
        const session = findSessionByPeer(peerId)
        if (!session) return
        session.history.push({ phaseId: data.phaseId, buttonId: data.buttonId, timestamp: Date.now() })
        session.currentPhaseId = data.nextPhaseId
        broadcastToSession(session, {
          type: 'state-change',
          state: getSessionState(session),
          action: { buttonId: data.buttonId, phaseId: data.phaseId },
        })
        break
      }

      case 'instructor-command': {
        const session = findSessionByPeer(peerId)
        if (!session) return
        const peerInfo = session.peers.get(peerId)
        if (peerInfo?.role !== 'instructor') return

        switch (data.command) {
          case 'pause':
            session.isPaused = true
            break
          case 'resume':
            session.isPaused = false
            break
          case 'restart':
            session.currentPhaseId = 'welcome'
            session.isPaused = false
            session.history = []
            session.startedAt = Date.now()
            break
          case 'goto':
            if (data.targetPhaseId) session.currentPhaseId = data.targetPhaseId
            break
        }
        broadcastToSession(session, { type: 'state-change', state: getSessionState(session) })
        break
      }

      case 'instructor-message': {
        const session = findSessionByPeer(peerId)
        if (!session) return
        broadcastToSession(session, {
          type: 'instructor-message',
          text: data.text,
          withRadioEffect: data.withRadioEffect ?? true,
        }, peerId)
        break
      }

      case 'subscribe-telemetry': {
        // Client sends their userId so bridge data gets routed to their session
        const session = findSessionByPeer(peerId)
        if (!session || !data.userId) return
        const peerInfo = session.peers.get(peerId)
        if (peerInfo) peerInfo.userId = data.userId
        session.telemetryUserIds.add(data.userId)
        userSessionMap.set(data.userId, session.code)
        break
      }

      case 'stick-input': {
        // Broadcast stick/throttle input to all peers in session (for PFD display)
        const session = findSessionByPeer(peerId)
        if (!session) return
        broadcastToSession(session, { type: 'stick-input', data: data.data }, peerId)
        break
      }
    }
  },

  close(peer) {
    const peerId = (peer as any).id ?? String(peer)
    for (const [code, session] of sessions) {
      const peerInfo = session.peers.get(peerId)
      if (peerInfo) {
        // Clean up telemetry subscription
        if (peerInfo.userId) {
          session.telemetryUserIds.delete(peerInfo.userId)
          userSessionMap.delete(peerInfo.userId)
        }
        session.peers.delete(peerId)
        if (peerInfo.role === 'participant') session.participantConnected = false
        if (peerInfo.role === 'instructor') session.instructorConnected = false
        broadcastToSession(session, { type: 'peer-left', role: peerInfo.role })
        // Clean up empty sessions
        if (session.peers.size === 0) sessions.delete(code)
        break
      }
    }
  },
})