lms-be/src/chat/chat.gateway.ts

import { Injectable } from '@nestjs/common';
import {
  ConnectedSocket,
  MessageBody,
  OnGatewayConnection,
  SubscribeMessage,
  WebSocketGateway,
  WebSocketServer,
} from '@nestjs/websockets';
import { Server, Socket } from 'socket.io';

import { AuthService, type AuthResponseUser } from '../auth/auth.service';
import { PrismaService } from '../prisma/prisma.service';

type ChatSocketUser = {
  id: string;
  email: string;
  role: string;
  roles?: string[];
};

export type ChatNotificationPayload = {
  type: 'chat.message.created';
  threadId: string;
  messageId: string;
  senderType: 'STUDENT' | 'AI' | 'INSTRUCTOR' | 'SYSTEM';
  preview: string;
  createdAt: string;
};

export type UserNotificationPayload =
  | ChatNotificationPayload
  | {
      type: 'platform.control.warning';
      sessionId: string;
      platformId: string | null;
      instructorPlatformId: string | null;
      endsAt: string;
      warningAt: string;
    }
  | {
      type: 'platform.control.granted';
      sessionId: string;
      startedAt: string;
      endsAt: string;
      warningAt: string;
    }
  | {
      type: 'platform.control.terminated';
      sessionId: string;
      platformId: string | null;
      instructorPlatformId: string | null;
      reason: string;
      message: string;
      endedAt: string;
    };

@Injectable()
@WebSocketGateway({
  cors: {
    origin: true,
    credentials: true,
  },
})
export class ChatGateway implements OnGatewayConnection {
  @WebSocketServer()
  private server!: Server;

  constructor(
    private readonly authService: AuthService,
    private readonly prisma: PrismaService,
  ) {}

  async handleConnection(client: Socket) {
    const accessToken = this.getAccessToken(client);
    if (!accessToken) {
      client.disconnect(true);
      return;
    }

    try {
      const { user } = await this.authService.me(accessToken);
      const lmsUser = await this.prisma.user.upsert({
        where: { email: user.email },
        update: this.toUserSyncData(user),
        create: {
          email: user.email,
          ...this.toUserSyncData(user),
        },
      });
      const socketUser: ChatSocketUser = {
        id: lmsUser.id,
        email: user.email,
        role: user.role,
        roles: user.roles,
      };

      client.data.user = socketUser;
      await client.join(this.userRoom(lmsUser.id));
      const organizationName = user.primaryAffiliation?.organizationName?.trim();
      if (organizationName) {
        await client.join(this.organizationRoom(organizationName));
      }

      if (this.isInstructor(socketUser)) {
        const activeCourses = await this.prisma.course.findMany({
          where: { isActive: true },
          select: { id: true },
        });
        await Promise.all(
          activeCourses.map((course) =>
            client.join(this.courseInstructorsRoom(course.id)),
          ),
        );
      }
    } catch {
      client.disconnect(true);
    }
  }

  @SubscribeMessage('chat.thread.join')
  async joinThread(
    @ConnectedSocket() client: Socket,
    @MessageBody() payload: { threadId?: string },
  ) {
    const threadId = payload.threadId;
    if (!threadId || !client.data.user) return { ok: false };

    await client.join(this.threadRoom(threadId));
    return { ok: true };
  }

  @SubscribeMessage('chat.thread.leave')
  async leaveThread(
    @ConnectedSocket() client: Socket,
    @MessageBody() payload: { threadId?: string },
  ) {
    const threadId = payload.threadId;
    if (!threadId) return { ok: false };

    await client.leave(this.threadRoom(threadId));
    return { ok: true };
  }

  notifyUser(userId: string, payload: UserNotificationPayload) {
    this.server.to(this.userRoom(userId)).emit(payload.type, payload);
  }

  notifyCourseInstructors(courseId: string, payload: ChatNotificationPayload) {
    this.server
      .to(this.courseInstructorsRoom(courseId))
      .emit(payload.type, payload);
  }

  notifyThread(threadId: string, payload: ChatNotificationPayload) {
    this.server.to(this.threadRoom(threadId)).emit(payload.type, payload);
  }

  notifyOrganization(organizationName: string, eventName: string, payload: unknown) {
    this.server.to(this.organizationRoom(organizationName)).emit(eventName, payload);
  }

  async notifyGroupMembers(groupId: string, eventName: string, payload: unknown) {
    const members = await this.prisma.studentGroupMember.findMany({
      where: { groupId },
      select: { userId: true },
    });

    members.forEach((member) => {
      this.server.to(this.userRoom(member.userId)).emit(eventName, payload);
    });
  }

  private getAccessToken(client: Socket) {
    const authToken = client.handshake.auth?.token;
    if (typeof authToken === 'string' && authToken) return authToken;

    const authorization = client.handshake.headers.authorization;
    if (!authorization) return null;

    const [type, token] = authorization.split(' ');
    return type === 'Bearer' ? token : null;
  }

  private isInstructor(user: ChatSocketUser) {
    const roles = new Set(
      [user.role, ...(user.roles ?? [])].map((role) => role.toUpperCase()),
    );
    return ['INSTRUCTOR', 'LMS_ADMIN'].some((role) => roles.has(role));
  }

  private toUserSyncData(user: AuthResponseUser) {
    return {
      authUserId: user.id,
      name: user.name,
      authRole: user.role,
      primaryOrganizationId: user.primaryAffiliation?.organizationId ?? null,
      primaryOrganizationName: user.primaryAffiliation?.organizationName ?? null,
      primaryOrganizationType: user.primaryAffiliation?.organizationType ?? null,
      department: user.primaryAffiliation?.department ?? null,
      position: user.primaryAffiliation?.position ?? null,
    };
  }

  private userRoom(userId: string) {
    return `user:${userId}`;
  }

  private courseInstructorsRoom(courseId: string) {
    return `course:${courseId}:instructors`;
  }

  private threadRoom(threadId: string) {
    return `thread:${threadId}`;
  }

  private organizationRoom(organizationName: string) {
    return `organization:${organizationName}`;
  }
}