import type { ServerWebSocket } from 'bun';
import type { Socket } from 'bun';
import * as sidecar from '@@/sidecar-registry';

type WSData = { userId: number; email: string; username: string; sessionId?: string };

type VncSession = {
  tcpSocket: Socket<{ ws: ServerWebSocket<WSData> }> | null;
  pendingMessages: Buffer[];
};

const sessions = new Map<ServerWebSocket<WSData>, VncSession>();

export const desktopWebsocket = {
  async open(ws: ServerWebSocket<WSData>) {
    const session: VncSession = { tcpSocket: null, pendingMessages: [] };
    sessions.set(ws, session);

    let port: number;
    try {
      const result = await sidecar.startVnc({
        email: ws.data.email,
        username: ws.data.username,
      });
      port = result.port;
    } catch (err) {
      const message = err instanceof Error ? err.message : 'Failed to start VNC session';
      console.error('[desktop] VNC start failed:', message);
      sessions.delete(ws);
      ws.close(4004, 'VNC server unavailable');
      return;
    }

    try {
      const tcpSocket = await Bun.connect({
        hostname: '127.0.0.1',
        port,
        socket: {
          data(_socket, data) {
            try {
              ws.sendBinary(Buffer.from(data));
            } catch {
              // ws already closed
            }
          },
          close() {
            sessions.delete(ws);
            try {
              ws.close();
            } catch {
              /* ignore */
            }
          },
          error(_socket, err) {
            console.error('[desktop] TCP error:', err.message);
            sessions.delete(ws);
            try {
              ws.close();
            } catch {
              /* ignore */
            }
          },
          connectError(_socket, err) {
            console.error('[desktop] TCP connect error:', err.message);
            sessions.delete(ws);
            try {
              ws.close(4004, 'VNC server unavailable');
            } catch {
              /* ignore */
            }
          },
          open(socket) {
            for (const msg of session.pendingMessages) {
              socket.write(msg);
            }
            session.pendingMessages = [];
          },
        },
        data: { ws },
      });

      session.tcpSocket = tcpSocket;
    } catch (err) {
      const message = err instanceof Error ? err.message : 'Failed to connect to VNC server';
      console.error('[desktop] VNC connect failed:', message);
      sessions.delete(ws);
      ws.close(4004, 'VNC server unavailable');
    }
  },

  message(ws: ServerWebSocket<WSData>, raw: string | Buffer) {
    const session = sessions.get(ws);
    if (!session) return;

    const data = typeof raw === 'string' ? Buffer.from(raw) : Buffer.from(raw);

    if (!session.tcpSocket) {
      session.pendingMessages.push(data);
      return;
    }

    try {
      session.tcpSocket.write(data);
    } catch {
      // ignore
    }
  },

  close(ws: ServerWebSocket<WSData>) {
    const session = sessions.get(ws);
    if (session?.tcpSocket) {
      try {
        session.tcpSocket.end();
      } catch {
        // ignore
      }
    }
    sessions.delete(ws);
  },

  drain() {},
};