import produce from 'immer';
import { atom, useSetAtom } from 'jotai';
import { MatrixClient, RoomMemberEvent, RoomMemberEventHandlerMap } from 'matrix-js-sdk';
import { useEffect } from 'react';

export const TYPING_TIMEOUT_MS = 5000; // 5 seconds

export type TypingReceipt = {
  userId: string;
  ts: number;
};
export type IRoomIdToTypingMembers = Map<string, TypingReceipt[]>;

type TypingMemberPutAction = {
  type: 'PUT';
  roomId: string;
  userId: string;
  ts: number;
};
type TypingMemberDeleteAction = {
  type: 'DELETE';
  roomId: string;
  userId: string;
};
export type IRoomIdToTypingMembersAction = TypingMemberPutAction | TypingMemberDeleteAction;

const baseRoomIdToTypingMembersAtom = atom<IRoomIdToTypingMembers>(new Map());

const putTypingMember = (
  roomToMembers: IRoomIdToTypingMembers,
  action: TypingMemberPutAction
): IRoomIdToTypingMembers => {
  let typingMembers = roomToMembers.get(action.roomId) ?? [];

  typingMembers = typingMembers.filter((receipt) => receipt.userId !== action.userId);
  typingMembers.push({
    userId: action.userId,
    ts: action.ts,
  });
  roomToMembers.set(action.roomId, typingMembers);
  return roomToMembers;
};

const deleteTypingMember = (
  roomToMembers: IRoomIdToTypingMembers,
  action: TypingMemberDeleteAction
): IRoomIdToTypingMembers => {
  let typingMembers = roomToMembers.get(action.roomId) ?? [];

  typingMembers = typingMembers.filter((receipt) => receipt.userId !== action.userId);
  if (typingMembers.length === 0) {
    roomToMembers.delete(action.roomId);
  } else {
    roomToMembers.set(action.roomId, typingMembers);
  }
  return roomToMembers;
};

const timeoutReceipt = (
  roomToMembers: IRoomIdToTypingMembers,
  roomId: string,
  userId: string,
  timeout: number
): boolean | undefined => {
  const typingMembers = roomToMembers.get(roomId) ?? [];

  const target = typingMembers.find((receipt) => receipt.userId === userId);
  if (!target) return undefined;

  return Date.now() - target.ts >= timeout;
};

export const roomIdToTypingMembersAtom = atom<
  IRoomIdToTypingMembers,
  [IRoomIdToTypingMembersAction],
  undefined
>(
  (get) => get(baseRoomIdToTypingMembersAtom),
  (get, set, action) => {
    const rToTyping = get(baseRoomIdToTypingMembersAtom);

    if (action.type === 'PUT') {
      set(
        baseRoomIdToTypingMembersAtom,
        produce(rToTyping, (draft) => putTypingMember(draft, action))
      );

      // remove typing receipt after some timeout
      // to prevent stuck typing members
      setTimeout(() => {
        const { roomId, userId } = action;
        const timeout = timeoutReceipt(
          get(baseRoomIdToTypingMembersAtom),
          roomId,
          userId,
          TYPING_TIMEOUT_MS
        );
        if (timeout) {
          set(
            baseRoomIdToTypingMembersAtom,
            produce(get(baseRoomIdToTypingMembersAtom), (draft) =>
              deleteTypingMember(draft, {
                type: 'DELETE',
                roomId,
                userId,
              })
            )
          );
        }
      }, TYPING_TIMEOUT_MS);
    }

    if (
      action.type === 'DELETE' &&
      rToTyping.get(action.roomId)?.find((receipt) => receipt.userId === action.userId)
    ) {
      set(
        baseRoomIdToTypingMembersAtom,
        produce(rToTyping, (draft) => deleteTypingMember(draft, action))
      );
    }
  }
);

export const useBindRoomIdToTypingMembersAtom = (
  mx: MatrixClient,
  typingMembersAtom: typeof roomIdToTypingMembersAtom
) => {
  const setTypingMembers = useSetAtom(typingMembersAtom);

  useEffect(() => {
    const handleTypingEvent: RoomMemberEventHandlerMap[RoomMemberEvent.Typing] = (
      event,
      member
    ) => {
      setTypingMembers({
        type: member.typing ? 'PUT' : 'DELETE',
        roomId: member.roomId,
        userId: member.userId,
        ts: Date.now(),
      });
    };

    mx.on(RoomMemberEvent.Typing, handleTypingEvent);
    return () => {
      mx.removeListener(RoomMemberEvent.Typing, handleTypingEvent);
    };
  }, [mx, setTypingMembers]);
};
