import {
  Injectable,
  Logger,
  OnModuleInit,
  UnauthorizedException,
} from '@nestjs/common';
import { HttpAdapterHost } from '@nestjs/core';
import type { FastifyRequest } from 'fastify';
import type { WebSocket } from 'ws';

import { NodesService } from './nodes.service';

@Injectable()
export class NodesGateway implements OnModuleInit {
  private readonly logger = new Logger(NodesGateway.name);

  constructor(
    private readonly adapterHost: HttpAdapterHost,
    private readonly nodesService: NodesService,
  ) {}

  onModuleInit(): void {
    const fastify = this.adapterHost.httpAdapter.getInstance();

    fastify.get(
      '/api/v1/nodes/ws',
      { websocket: true },
      (socket: WebSocket, request: FastifyRequest) => {
        void this.handleConnection(socket, request);
      },
    );
  }

  private async handleConnection(
    socket: WebSocket,
    request: FastifyRequest,
  ): Promise<void> {
    const query = request.query as Record<string, string | string[] | undefined>;
    const nodeId = this.readQueryParam(query['nodeId']);
    const token = this.readQueryParam(query['token']);

    if (!nodeId || !token) {
      socket.close(4400, 'nodeId and token are required');
      return;
    }

    try {
      await this.nodesService.validateEnrollment(nodeId, token);
    } catch (error) {
      const reason =
        error instanceof UnauthorizedException
          ? 'Invalid enrollment credentials'
          : 'Enrollment validation failed';
      socket.close(4401, reason);
      return;
    }

    this.nodesService.registerConnection(nodeId, socket);

    socket.on('close', () => {
      this.nodesService.unregisterConnection(nodeId, socket);
    });

    socket.on('message', (data) => {
      void this.nodesService.handleAgentMessage(nodeId, data.toString());
    });

    socket.on('error', (error) => {
      this.logger.error(`WebSocket error for node ${nodeId}: ${error.message}`);
    });

    socket.send(JSON.stringify(this.nodesService.createHelloAck(nodeId)));
  }

  private readQueryParam(
    value: string | string[] | undefined,
  ): string | undefined {
    if (Array.isArray(value)) {
      return value[0];
    }

    return value;
  }
}