corrosion-admin-panel/backend-nest/src/modules/console/console.gateway.ts

import {
  WebSocketGateway,
  WebSocketServer,
  SubscribeMessage,
  OnGatewayConnection,
  OnGatewayDisconnect,
  MessageBody,
  ConnectedSocket,
} from '@nestjs/websockets';
import WebSocket, { Server } from 'ws';
import { IncomingMessage } from 'http';
import { Logger, UnauthorizedException } from '@nestjs/common';
import { JwtService } from '@nestjs/jwt';
import { NatsService } from '../../services/nats.service';

interface ClientMeta {
  licenseId: string;
  userId: string;
}

/**
 * Console Gateway
 *
 * NOTE: This gateway is NOT currently loaded (ConsoleModule not imported in AppModule).
 * Console I/O is handled by NatsBridgeGateway instead.
 * Kept for potential future use as a dedicated console-only WebSocket endpoint.
 */
@WebSocketGateway({ path: '/api/console-ws' })
export class ConsoleGateway implements OnGatewayConnection, OnGatewayDisconnect {
  @WebSocketServer()
  server: Server;

  private readonly logger = new Logger(ConsoleGateway.name);
  private clientMeta = new Map<WebSocket, ClientMeta>();
  private licenseClients = new Map<string, Set<WebSocket>>();

  constructor(
    private readonly jwtService: JwtService,
    private readonly natsService: NatsService,
  ) {}

  async handleConnection(client: WebSocket, request: IncomingMessage) {
    try {
      const url = new URL(request.url || '/', `http://${request.headers.host}`);
      const token = url.searchParams.get('token');

      if (!token) {
        throw new UnauthorizedException('No token provided');
      }

      const payload = this.jwtService.verify(token);
      const licenseId = payload.license_id;

      if (!licenseId) {
        throw new UnauthorizedException('Invalid token: no license_id');
      }

      this.clientMeta.set(client, { licenseId, userId: payload.sub });

      if (!this.licenseClients.has(licenseId)) {
        this.licenseClients.set(licenseId, new Set());
      }
      this.licenseClients.get(licenseId)!.add(client);

      this.logger.log(`Client connected to license ${licenseId}`);
    } catch (error) {
      const message = error instanceof Error ? error.message : 'Unknown error';
      this.logger.error(`Connection failed: ${message}`);
      client.close(4001, message);
    }
  }

  handleDisconnect(client: WebSocket) {
    const meta = this.clientMeta.get(client);
    if (meta?.licenseId) {
      this.licenseClients.get(meta.licenseId)?.delete(client);
      if (this.licenseClients.get(meta.licenseId)?.size === 0) {
        this.licenseClients.delete(meta.licenseId);
      }
    }
    this.clientMeta.delete(client);
  }

  @SubscribeMessage('console_input')
  async handleConsoleInput(
    @ConnectedSocket() client: WebSocket,
    @MessageBody() data: { command: string },
  ) {
    const meta = this.clientMeta.get(client);
    if (!meta?.licenseId) return;

    if (!data.command) {
      return { error: 'Command is required' };
    }

    this.logger.debug(`Console input from ${meta.licenseId}: ${data.command}`);

    await this.natsService.sendServerCommand(meta.licenseId, 'command', {
      command: data.command,
    });

    return { success: true };
  }

  sendToLicense(licenseId: string, event: string, data: any) {
    const clients = this.licenseClients.get(licenseId);
    if (!clients) return;

    const message = JSON.stringify({ event, data });
    for (const client of clients) {
      if (client.readyState === WebSocket.OPEN) {
        client.send(message);
      }
    }
  }

  broadcastConsoleOutput(licenseId: string, output: string) {
    this.sendToLicense(licenseId, 'console_output', { output });
  }
}