#!/usr/bin/env node
/**
 * Simple y-websocket server with connection logging and JWT authentication.
 * Run with: node server.mjs
 */

import http from 'http'
import crypto from 'node:crypto'
import { WebSocketServer } from 'ws'
import * as Y from 'yjs'
import * as syncProtocol from 'y-protocols/sync'
import * as awarenessProtocol from 'y-protocols/awareness'
import * as encoding from 'lib0/encoding'
import * as decoding from 'lib0/decoding'

const PORT = process.env.PORT || 3002
const HOST = process.env.HOST || '0.0.0.0'
const WINDMILL_BASE_URL = process.env.WINDMILL_BASE_URL || process.env.BASE_INTERNAL_URL
const REQUIRE_SIGNED_REQUESTS = process.env.REQUIRE_SIGNED_MULTIPLAYER_REQUESTS !== 'false'

const messageSync = 0
const messageAwareness = 1

// Caps on what an unauthenticated peer may buffer while its token is verified.
// A real client only has sync step 1 and its first awareness update in flight
// there, so both limits are far above legitimate use.
const MAX_PREAUTH_MESSAGES = 32
const MAX_PREAUTH_BYTES = 1024 * 1024

// --- JWT verification ---

let cachedPublicKey = null
let publicKeyFetchPromise = null

function base64urlDecode(str) {
  const padding = '='.repeat((4 - str.length % 4) % 4)
  const base64 = str.replace(/-/g, '+').replace(/_/g, '/') + padding
  const binary = atob(base64)
  return Uint8Array.from(binary, c => c.charCodeAt(0))
}

async function getPublicKey() {
  if (cachedPublicKey) return cachedPublicKey
  if (publicKeyFetchPromise) return publicKeyFetchPromise

  if (!WINDMILL_BASE_URL) {
    console.warn(`[${new Date().toISOString()}] WINDMILL_BASE_URL not set - cannot fetch public key`)
    return null
  }

  publicKeyFetchPromise = (async () => {
    try {
      const jwksUrl = `${WINDMILL_BASE_URL.replace(/\/$/, '')}/api/debug/jwks`
      console.log(`[${new Date().toISOString()}] Fetching JWKS from ${jwksUrl}`)
      const response = await fetch(jwksUrl)
      if (!response.ok) {
        throw new Error(`Failed to fetch JWKS: ${response.status} ${response.statusText}`)
      }

      const jwks = await response.json()
      if (!jwks.keys || jwks.keys.length === 0) {
        throw new Error('No keys in JWKS')
      }

      const jwk = jwks.keys[0]
      if (jwk.kty !== 'OKP' || jwk.crv !== 'Ed25519') {
        throw new Error(`Unsupported key type: ${jwk.kty}/${jwk.crv}`)
      }

      const publicKeyBytes = base64urlDecode(jwk.x)
      const key = await crypto.subtle.importKey(
        'raw',
        publicKeyBytes,
        { name: 'Ed25519' },
        true,
        ['verify']
      )

      cachedPublicKey = key
      console.log(`[${new Date().toISOString()}] Successfully loaded Ed25519 public key from JWKS`)
      return key
    } catch (error) {
      console.error(`[${new Date().toISOString()}] Failed to fetch/parse JWKS: ${error}`)
      return null
    } finally {
      publicKeyFetchPromise = null
    }
  })()

  return publicKeyFetchPromise
}

/**
 * Verify a JWT multiplayer token.
 * Returns null if valid, or an error message if invalid.
 */
async function verifyToken(token, docName) {
  const publicKey = await getPublicKey()

  if (!publicKey) {
    if (REQUIRE_SIGNED_REQUESTS) {
      return 'Public key not available but signed requests are required. Set WINDMILL_BASE_URL.'
    }
    console.warn(`[${new Date().toISOString()}] Public key not available - signature verification disabled`)
    return null
  }

  const parts = token.split('.')
  if (parts.length !== 3) {
    return 'Invalid JWT format'
  }

  const [headerB64, claimsB64, signatureB64] = parts

  try {
    const message = new TextEncoder().encode(`${headerB64}.${claimsB64}`)
    const signature = base64urlDecode(signatureB64)

    const isValid = await crypto.subtle.verify(
      { name: 'Ed25519' },
      publicKey,
      signature,
      message
    )

    if (!isValid) {
      return 'Invalid JWT signature'
    }

    const claimsJson = new TextDecoder().decode(base64urlDecode(claimsB64))
    const claims = JSON.parse(claimsJson)

    // Check expiration
    const now = Math.floor(Date.now() / 1000)
    if (now > claims.exp) {
      return `Token expired: ${now - claims.exp} seconds ago`
    }

    // Check purpose
    if (claims.purpose !== 'multiplayer') {
      return `Invalid token purpose: ${claims.purpose}`
    }

    // Check workspace matches the doc room (format: "{workspace}/{path}")
    const docWorkspace = docName.split('/')[0]
    if (docWorkspace && claims.workspace_id !== docWorkspace) {
      return `Token workspace "${claims.workspace_id}" does not match room workspace "${docWorkspace}"`
    }

    return null
  } catch (error) {
    return `JWT verification error: ${error}`
  }
}

// --- Y.js document management ---

// Store docs in memory
const docs = new Map()

const getYDoc = (docname) => {
  let doc = docs.get(docname)
  if (!doc) {
    doc = new Y.Doc()
    doc.name = docname
    docs.set(docname, doc)
  }
  return doc
}

/**
 * Every error message is peer-controlled: `applyAwarenessUpdate` runs
 * `JSON.parse` on the peer's bytes and V8 quotes the offending input back. Only
 * printable ASCII survives, so nothing in it can forge or flood a log line.
 */
const describeError = (error) => String(error?.message ?? error).replace(/[^\x20-\x7e]+/g, ' ').slice(0, 200)

const send = (conn, message) => {
  if (conn.readyState === 1) { // WebSocket.OPEN
    conn.send(message, err => { if (err) console.error(err) })
  }
}

const setupWSConnection = (conn, req, docName, bufferedMessages = []) => {
  const doc = getYDoc(docName)
  const clientIp = req.socket.remoteAddress

  // Initialize awareness
  if (!doc.awareness) {
    doc.awareness = new awarenessProtocol.Awareness(doc)
  }

  const awareness = doc.awareness

  // Track connections per doc
  if (!doc.conns) doc.conns = new Set()
  doc.conns.add(conn)

  const messageHandler = (message) => {
    // A frame that arrived before we decided to close this connection (or that a
    // replay below already made moot) must not still be applied to the document.
    if (conn.readyState !== 1) return // WebSocket.OPEN

    try {
      const data = new Uint8Array(message)
      const decoder = decoding.createDecoder(data)
      const messageType = decoding.readVarUint(decoder)

      switch (messageType) {
        case messageSync: {
          const encoder = encoding.createEncoder()
          encoding.writeVarUint(encoder, messageSync)
          syncProtocol.readSyncMessage(decoder, encoder, doc, null)
          if (encoding.length(encoder) > 1) {
            send(conn, encoding.toUint8Array(encoder))
          }
          break
        }
        case messageAwareness:
          awarenessProtocol.applyAwarenessUpdate(awareness, decoding.readVarUint8Array(decoder), conn)
          break
      }
    } catch (error) {
      // The decoders throw on unparseable input and `ws` re-emits a listener's
      // exception on the process, so one bad frame from one peer would end the
      // server for every document and client. Drop only the offender, with RFC
      // 6455's 1007; the frame itself is untrusted and is never logged.
      console.warn(`[${new Date().toISOString()}] MALFORMED MESSAGE: doc="${docName}" from=${clientIp} error="${describeError(error)}"`)
      conn.close(1007, 'Invalid message')
    }
  }
  conn.on('message', messageHandler)

  // Send initial sync step 1
  {
    const encoder = encoding.createEncoder()
    encoding.writeVarUint(encoder, messageSync)
    syncProtocol.writeSyncStep1(encoder, doc)
    send(conn, encoding.toUint8Array(encoder))
  }

  // Send awareness states
  const awarenessStates = awareness.getStates()
  if (awarenessStates.size > 0) {
    const encoder = encoding.createEncoder()
    encoding.writeVarUint(encoder, messageAwareness)
    encoding.writeVarUint8Array(encoder, awarenessProtocol.encodeAwarenessUpdate(awareness, Array.from(awarenessStates.keys())))
    send(conn, encoding.toUint8Array(encoder))
  }

  // Broadcast awareness changes
  const awarenessChangeHandler = ({ added, updated, removed }, origin) => {
    const changedClients = added.concat(updated).concat(removed)
    const encoder = encoding.createEncoder()
    encoding.writeVarUint(encoder, messageAwareness)
    encoding.writeVarUint8Array(encoder, awarenessProtocol.encodeAwarenessUpdate(awareness, changedClients))
    const message = encoding.toUint8Array(encoder)
    doc.conns.forEach(c => send(c, message))
  }
  awareness.on('update', awarenessChangeHandler)

  // Broadcast doc updates
  const updateHandler = (update, origin) => {
    const encoder = encoding.createEncoder()
    encoding.writeVarUint(encoder, messageSync)
    syncProtocol.writeUpdate(encoder, update)
    const message = encoding.toUint8Array(encoder)
    doc.conns.forEach(c => {
      if (c !== origin) send(c, message)
    })
  }
  doc.on('update', updateHandler)

  conn.on('close', () => {
    doc.conns.delete(conn)
    awareness.off('update', awarenessChangeHandler)
    doc.off('update', updateHandler)

    // Clean up awareness for this connection
    awarenessProtocol.removeAwarenessStates(awareness, [doc.clientID], null)
  })

  // Replay, in order, the messages that arrived while the token was being
  // verified, now that the handlers above are in place. The log line is the only
  // outside evidence that a frame took this path rather than the live handler,
  // and marks a client that beat the JWKS fetch on a slow-starting instance.
  if (bufferedMessages.length > 0) {
    console.log(`[${new Date().toISOString()}] REPLAY: doc="${docName}" from=${clientIp} messages=${bufferedMessages.length}`)
  }
  bufferedMessages.forEach(messageHandler)
}

// --- HTTP + WebSocket server ---

const server = http.createServer((req, res) => {
  // Strip /ws_mp/ prefix if present (when accessed without reverse proxy path stripping)
  if (req.url?.startsWith('/ws_mp/')) {
    req.url = req.url.slice('/ws_mp'.length)
  } else if (req.url === '/ws_mp') {
    req.url = '/'
  }
  console.log(`[${new Date().toISOString()}] HTTP ${req.method} ${req.url} from=${req.socket.remoteAddress}`)
  if (req.url === '/' || req.url === '/health') {
    res.writeHead(200, {
      'Content-Type': 'application/json',
      'Access-Control-Allow-Origin': '*'
    })
    res.end(JSON.stringify({ status: 'ok', service: 'multiplayer' }))
  } else {
    res.writeHead(404)
    res.end('not found')
  }
})

const wss = new WebSocketServer({ server })

// `ws` forwards the HTTP server's errors here, and those are fatal: a failed
// listen leaves nothing to serve, and exiting 0 would read as a clean shutdown.
// `process.exitCode` rather than `process.exit()`, which truncates this line.
wss.on('error', (error) => {
  console.error(`[${new Date().toISOString()}] WEBSOCKET SERVER ERROR: ${describeError(error)}`)
  process.exitCode = 1
})

wss.on('connection', async (ws, req) => {
  let docName = req.url?.slice(1).split('?')[0] || 'unknown'

  // Strip ws_mp/ prefix if present (when accessed without reverse proxy path stripping)
  if (docName.startsWith('ws_mp/')) {
    docName = docName.slice('ws_mp/'.length)
  }

  const clientIp = req.socket.remoteAddress

  // A frame `ws` cannot parse at the protocol level fails in its Receiver, never
  // reaching the handler below, and an unhandled 'error' on an EventEmitter ends
  // the process. Attached before authentication, since the pre-auth window is
  // exposed too; `ws` has already closed the connection by the time this runs.
  ws.on('error', (error) => {
    console.warn(`[${new Date().toISOString()}] SOCKET ERROR: doc="${docName}" from=${clientIp} error="${describeError(error)}"`)
  })

  // Handle ping test — respond and close immediately
  if (docName === '__ping__') {
    console.log(`[${new Date().toISOString()}] WS ping from=${clientIp}`)
    ws.send(JSON.stringify({ type: 'pong', service: 'multiplayer' }))
    ws.close()
    return
  }

  // A y-websocket client sends sync step 1 as soon as the socket opens, which can
  // be before token verification resolves (the first connection after startup has
  // to wait for the JWKS fetch). `ws` drops messages emitted with no listener
  // attached, so buffer them here and replay them once the connection is accepted.
  // The peer is not authenticated yet and the JWKS fetch has no timeout, so what
  // it may buffer is capped.
  const bufferedMessages = []
  let bufferedBytes = 0
  const bufferMessage = (message) => {
    bufferedBytes += message.length
    if (bufferedMessages.length >= MAX_PREAUTH_MESSAGES || bufferedBytes > MAX_PREAUTH_BYTES) {
      console.warn(`[${new Date().toISOString()}] REJECTED: doc="${docName}" from=${clientIp} reason="too much data before authentication"`)
      bufferedMessages.length = 0
      ws.off('message', bufferMessage)
      ws.close(1009, 'Too much data before authentication')
      return
    }
    bufferedMessages.push(message)
  }
  ws.on('message', bufferMessage)

  // Verify JWT token
  const urlParams = new URLSearchParams(req.url?.split('?')[1] || '')
  const token = urlParams.get('token')

  if (!token) {
    if (REQUIRE_SIGNED_REQUESTS) {
      console.warn(`[${new Date().toISOString()}] REJECTED: doc="${docName}" from=${clientIp} reason="no token"`)
      ws.off('message', bufferMessage)
      ws.close(4401, 'Authentication required')
      return
    }
    console.warn(`[${new Date().toISOString()}] WARN: no token for doc="${docName}" from=${clientIp} (signed requests not required)`)
  } else {
    const error = await verifyToken(token, docName)
    if (error) {
      console.warn(`[${new Date().toISOString()}] REJECTED: doc="${docName}" from=${clientIp} reason="${error}"`)
      ws.off('message', bufferMessage)
      ws.close(4403, 'Token verification failed')
      return
    }
  }

  ws.off('message', bufferMessage)

  // The socket may already be gone: closed by the peer while the token was being
  // verified, or by the buffer cap above. Setting a doc connection up on it would
  // add it to `doc.conns` with a 'close' listener that can no longer fire.
  if (ws.readyState !== 1) { // WebSocket.OPEN
    return
  }

  console.log(`[${new Date().toISOString()}] CONNECT: doc="${docName}" from=${clientIp}`)

  ws.on('close', () => {
    console.log(`[${new Date().toISOString()}] DISCONNECT: doc="${docName}" from=${clientIp}`)
  })

  setupWSConnection(ws, req, docName, bufferedMessages)
})

server.listen(PORT, HOST, () => {
  console.log(`[${new Date().toISOString()}] Multiplayer server running at ${HOST}:${PORT}`)
  if (REQUIRE_SIGNED_REQUESTS) {
    console.log(`[${new Date().toISOString()}] Signed requests REQUIRED (set REQUIRE_SIGNED_MULTIPLAYER_REQUESTS=false to disable)`)
  } else {
    console.log(`[${new Date().toISOString()}] Signed requests DISABLED`)
  }

  // Warm the JWKS cache so the first connection does not have to wait for it.
  // Best-effort insurance only: getPublicKey() stays the lazy fallback, logs its
  // own failures and resolves to null rather than rejecting, and a connection
  // arriving before this resolves is handled by the message buffer.
  if (WINDMILL_BASE_URL) {
    getPublicKey()
  }
})