import { WebSocketServer } from 'ws'
import { setupWebSocketHandler } from './src/websocket-handler.js'
import { printServerInfo } from './src/server-utils.js'
import { loadState, saveState } from './src/persist.js'

export default function websocketPlugin() {
  return {
    name: 'vite-plugin-websocket',
    configureServer(server) {
      const wss = new WebSocketServer({ noServer: true })
      setupWebSocketHandler(wss, { initialState: loadState(), onStateChange: saveState })

      // Rewrite /display → / (index.html) e /controller → /controller.html
      server.middlewares.use((req, _res, next) => {
        if (req.url === '/display' || req.url === '/display/') req.url = '/'
        else if (req.url === '/controller' || req.url === '/controller/') req.url = '/controller.html'
        next()
      })

      server.httpServer.on('upgrade', (request, socket, head) => {
        const pathname = new URL(request.url, `http://${request.headers.host}`).pathname
        if (pathname === '/ws') {
          wss.handleUpgrade(request, socket, head, (ws) => {
            wss.emit('connection', ws, request)
          })
        }
      })

      server.httpServer.once('listening', () => {
        const { port } = server.httpServer.address()
        setTimeout(() => printServerInfo(port), 100)
      })
    }
  }
}