websocket-server.ts

import WebSocket, { WebSocketServer } from "ws"
import http from "http"
import fetch from "node-fetch" // For OpenSky API calls

// =============================================================================
// WebSocket Server for Real-time Flight Data
// =============================================================================
// This server provides real-time flight data (either mock or from OpenSky Network)
// to connected WebSocket clients.
//
// How to run:
// 1. Ensure 'ws' and 'node-fetch' are installed (pnpm add ws node-fetch @types/ws)
// 2. Compile: `npx tsc websocket-server.ts` (if needed, or use ts-node)
// 3. Run: `node websocket-server.js` or `npx ts-node websocket-server.ts`
//
// The server listens on the port specified by WS_PORT environment variable,
// or defaults to 8080.
// =============================================================================

// --- Copied and adapted from app/api/flights/route.ts ---
interface ProcessedFlight {
  id: string
  callsign: string
  latitude: number
  longitude: number
  altitude: number
  speed: number
  heading: number
  status: "En Route" | "Landing" | "Takeoff" | "Taxiing" | "Boarding" | "Delayed"
  aircraft_type: string
  origin: string
  destination: string
  squawk: string
  registration: string
  country: string
  lastContact: number
  onGround: boolean
  verticalRate: number
}

function generateEnhancedMockFlightData(limit = 200): ProcessedFlight[] {
  const airlines = [
    { code: "UAL", name: "United Airlines", country: "United States" },
    { code: "DAL", name: "Delta Air Lines", country: "United States" },
    { code: "AAL", name: "American Airlines", country: "United States" },
    { code: "SWA", name: "Southwest Airlines", country: "United States" },
    { code: "JBU", name: "JetBlue Airways", country: "United States" },
    { code: "BAW", name: "British Airways", country: "United Kingdom" },
    { code: "AFR", name: "Air France", country: "France" },
    { code: "DLH", name: "Lufthansa", country: "Germany" },
  ]
  const aircraftTypes = [
    { type: "B737-800", cruiseAlt: 37000, cruiseSpeed: 450 },
    { type: "A320-200", cruiseAlt: 36000, cruiseSpeed: 440 },
    { type: "B777-300ER", cruiseAlt: 41000, cruiseSpeed: 490 },
    { type: "A350-900", cruiseAlt: 42000, cruiseSpeed: 485 },
  ]
  const airports = [
    { code: "LAX", lat: 33.9425, lng: -118.4081, city: "Los Angeles" },
    { code: "JFK", lat: 40.6413, lng: -73.7781, city: "New York" },
    { code: "LHR", lat: 51.47, lng: -0.4543, city: "London" },
    { code: "CDG", lat: 49.0097, lng: 2.5479, city: "Paris" },
    { code: "NRT", lat: 35.772, lng: 140.3929, city: "Tokyo" },
  ]
  const flights: ProcessedFlight[] = []
  const currentTime = Math.floor(Date.now() / 1000)
  for (let i = 0; i < Math.min(limit, 200); i++) {
    const airline = airlines[Math.floor(Math.random() * airlines.length)]
    const aircraft = aircraftTypes[Math.floor(Math.random() * aircraftTypes.length)]
    const origin = airports[Math.floor(Math.random() * airports.length)]
    let destination = airports[Math.floor(Math.random() * airports.length)]
    while (destination.code === origin.code) {
      destination = airports[Math.floor(Math.random() * airports.length)]
    }
    const progress = Math.random()
    const lat = origin.lat + (destination.lat - origin.lat) * progress + (Math.random() - 0.5) * 2
    const lng = origin.lng + (destination.lng - origin.lng) * progress + (Math.random() - 0.5) * 2
    let altitude: number, speed: number, status: ProcessedFlight["status"], onGround: boolean
    if (progress < 0.05) {
      altitude = Math.random() * 5000 + 500
      speed = Math.random() * 100 + 150
      status = "Takeoff"
      onGround = true
    } else if (progress > 0.95) {
      altitude = Math.random() * 3000 + 200
      speed = Math.random() * 80 + 120
      status = "Landing"
      onGround = false // Approaching
    } else {
      altitude = aircraft.cruiseAlt + (Math.random() - 0.5) * 4000
      speed = aircraft.cruiseSpeed + (Math.random() - 0.5) * 50
      status = "En Route"
      onGround = false
    }
    const deltaLng = destination.lng - lng
    const deltaLat = destination.lat - lat
    let heading = (Math.atan2(deltaLng, deltaLat) * 180) / Math.PI
    if (heading < 0) heading += 360
    const flightNumber = Math.floor(Math.random() * 9000) + 1000
    const callsign = `${airline.code}${flightNumber}`
    flights.push({
      id: `mock_ws_${i.toString().padStart(4, "0")}`,
      callsign,
      latitude: parseFloat(lat.toFixed(6)),
      longitude: parseFloat(lng.toFixed(6)),
      altitude: Math.max(0, parseFloat(altitude.toFixed(0))),
      speed: Math.max(0, parseFloat(speed.toFixed(0))),
      heading: parseFloat(heading.toFixed(0)),
      status,
      aircraft_type: aircraft.type,
      origin: origin.code,
      destination: destination.code,
      squawk: Math.floor(Math.random() * 7777).toString().padStart(4, "0"),
      registration: `N${Math.floor(Math.random() * 900)}${String.fromCharCode(65 + Math.floor(Math.random() * 26))}${String.fromCharCode(65 + Math.floor(Math.random() * 26))}`,
      country: airline.country,
      lastContact: currentTime - Math.floor(Math.random() * 30),
      onGround,
      verticalRate: onGround ? 0 : parseFloat(((Math.random() - 0.5) * 1000).toFixed(0)),
    })
  }
  return flights
}
// --- End Copied ---

const PORT = process.env.WS_PORT || 8080
const server = http.createServer() // Basic HTTP server to host the WebSocket server
const wss = new WebSocketServer({ server })

// Configuration Constants
const FETCH_INTERVAL = 10000 // Interval for fetching/broadcasting data (milliseconds)
const OPENSKY_TIMEOUT = 8000 // Timeout for OpenSky API requests (milliseconds)
const MAX_FLIGHTS_TO_BROADCAST = 200 // Max number of flight records to send to clients per update

// Server state
let useMockData = true // Default to mock data. Can be changed by client messages.

console.log(`[WS Server] WebSocket server starting on port ${PORT}`)

async function fetchOpenSkyData(): Promise<ProcessedFlight[] | null> {
  console.log("[WS Server] Attempting OpenSky Network API...")
  try {
    // @ts-ignore
    const response = await fetch("https://opensky-network.org/api/states/all", {
      method: "GET",
      headers: {
        "User-Agent": "Aircraft-Tracking-System-WS/1.0",
        Accept: "application/json",
      },
      timeout: OPENSKY_TIMEOUT,
    })

    if (response.ok) {
      const data = await response.json()
      if (data.states && Array.isArray(data.states) && data.states.length > 0) {
        console.log(`[WS Server] OpenSky success: ${data.states.length} states received`)
        const processedFlights: ProcessedFlight[] = data.states
          .filter((state: any[]) => state && state.length >= 17 && state[5] !== null && state[6] !== null && !isNaN(Number(state[5])) && !isNaN(Number(state[6])))
          .slice(0, MAX_FLIGHTS_TO_BROADCAST * 2) // Fetch more initially, then limit after processing
          .map((state: any[], index: number): ProcessedFlight => {
            const [icao24, callsign, origin_country, , last_contact, longitude, latitude, baro_altitude, on_ground, velocity, true_track, vertical_rate] = state
            const safeNumber = (value: any, fallback = 0): number => {
              const num = Number(value)
              return isFinite(num) && !isNaN(num) ? num : fallback
            }
            return {
              id: icao24 || `live_ws_${index}`,
              callsign: (callsign?.trim() || `UNKN${index.toString().padStart(3, "0")}`).substring(0, 10),
              latitude: safeNumber(latitude),
              longitude: safeNumber(longitude),
              altitude: Math.max(0, safeNumber(baro_altitude)),
              speed: Math.max(0, safeNumber(velocity) * 1.94384), // m/s to knots
              heading: safeNumber(true_track, Math.random() * 360),
              status: on_ground ? "Taxiing" : "En Route",
              aircraft_type: "Unknown",
              origin: "LIVE",
              destination: "DATA",
              squawk: "0000",
              registration: (icao24 || `REG_WS_${index}`).toUpperCase().substring(0, 10),
              country: (origin_country || "Unknown").substring(0, 50),
              lastContact: safeNumber(last_contact, Math.floor(Date.now() / 1000)),
              onGround: Boolean(on_ground),
              verticalRate: safeNumber(vertical_rate),
            }
          })
        return processedFlights.slice(0, MAX_FLIGHTS_TO_BROADCAST)
      }
    }
    console.error(`[WS Server] OpenSky API returned ${response.status}: ${response.statusText}`)
    return null
  } catch (apiError) {
    console.error(`[WS Server] OpenSky fetch error:`, apiError)
    return null
  }
}

/**
 * Fetches flight data, trying OpenSky first (if not in mock mode)
 * and falling back to generated mock data.
 * @returns {Promise<FlightDataResponse>} The flight data response to be broadcast.
 */
async function getFlightData(): Promise<FlightDataResponse> {
  let flights: ProcessedFlight[] = []
  let source = "mock" // Default source
  let warning: string | undefined = undefined

  if (!useMockData) {
    console.log("[WS Server] Attempting to fetch live data from OpenSky.")
    const liveFlights = await fetchOpenSkyData()
    if (liveFlights && liveFlights.length > 0) {
      flights = liveFlights
      source = "opensky"
      console.log(`[WS Server] Successfully fetched ${flights.length} live flights from OpenSky.`)
    } else {
      // Fallback to mock data if OpenSky fails or returns no data
      flights = generateEnhancedMockFlightData(MAX_FLIGHTS_TO_BROADCAST)
      source = "fallback"
      warning = "OpenSky Network API is currently unavailable or returned no data. Using fallback mock flight data."
      console.warn(`[WS Server] OpenSky unavailable or no data. Using ${flights.length} fallback mock flights. Warning: ${warning}`)
    }
  } else {
    // Use mock data if explicitly requested
    flights = generateEnhancedMockFlightData(MAX_FLIGHTS_TO_BROADCAST)
    source = "mock"
    console.log(`[WS Server] Using ${flights.length} generated mock flights as requested.`)
  }

  return {
    flights,
    timestamp: Date.now(),
    source,
    warning,
    total: flights.length, // total is the number of flights being sent
    processed: flights.length, // processed is also the number of flights being sent
  }
}

interface FlightDataResponse {
  flights: ProcessedFlight[]
  timestamp: number
  source: string
  warning?: string
  total?: number
  processed?: number
}

function broadcastFlightData() {
  if (wss.clients.size === 0) {
    // console.log("[WS Server] No clients connected, skipping broadcast.")
    return
  }

  getFlightData()
    .then((data) => {
      const jsonData = JSON.stringify(data)
      wss.clients.forEach((client) => {
        if (client.readyState === WebSocket.OPEN) {
          client.send(jsonData)
        }
      })
      console.log(`[WS Server] Broadcasted ${data.flights.length} flights from ${data.source} to ${wss.clients.size} clients.`)
      if (data.warning) console.warn(`[WS Server] Broadcast warning: ${data.warning}`)
    })
    .catch((error) => {
      console.error("[WS Server] Error getting flight data for broadcast:", error)
    })
}

wss.on("connection", (ws) => {
  console.log("[WS Server] Client connected.")

  // Send initial data immediately upon connection
  getFlightData().then(data => {
    if (ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify(data))
      console.log("[WS Server] Sent initial data to newly connected client.")
    }
  }).catch(error => {
    console.error("[WS Server] Error preparing initial data for client:", error)
  })

  // Handle messages from clients
  // Expected message format: { type: "SET_MOCK_DATA", payload: boolean }
  ws.on("message", (message) => {
    try {
      const parsedMessage = JSON.parse(message.toString())
      console.log("[WS Server] Received message from client:", parsedMessage)

      if (parsedMessage.type === "SET_MOCK_DATA" && typeof parsedMessage.payload === "boolean") {
        const oldMockStatus = useMockData
        useMockData = parsedMessage.payload
        console.log(`[WS Server] Client updated 'useMockData' from ${oldMockStatus} to ${useMockData}.`)
        // Immediately broadcast new data based on the updated preference
        broadcastFlightData()
      } else {
        console.warn("[WS Server] Received unknown message format from client:", parsedMessage)
      }
    } catch (e) {
      console.error("[WS Server] Failed to parse message from client or invalid message structure:", message.toString(), e)
    }
  })

  ws.on("close", () => {
    console.log("[WS Server] Client disconnected.")
  })

  ws.on("error", (error) => {
    console.error("[WS Server] WebSocket error:", error)
  })
})

server.listen(PORT, () => {
  console.log(`[WS Server] HTTP server for WebSocket is running on http://localhost:${PORT}`)
})

// Start broadcasting periodically
setInterval(broadcastFlightData, FETCH_INTERVAL)

console.log(`[WS Server] Initialized. Broadcasting data every ${FETCH_INTERVAL / 1000} seconds.`)

// Graceful shutdown
process.on('SIGINT', () => {
  console.log('[WS Server] Shutting down...');
  wss.clients.forEach(client => {
    client.close();
  });
  server.close(() => {
    console.log('[WS Server] Shutdown complete.');
    process.exit(0);
  });
});