cozsweet-frontend-nextjs/src/stores/chat/machine/unlock-flow.ts

import { type DoneActorEvent } from "xstate";

import { applyPromotionUnlockOutput } from "../helper/promotion";
import {
  applySingleUnlockOutput,
  countLockedHistoryMessages,
  type UnlockMessageOutput,
} from "../helper/unlock";
import type { UnlockHistoryOutput } from "./actors/unlock";
import { sendMachineSetup } from "./send-flow";
import { createChatActorActionSetup } from "./setup";
import { getCharacterErrorCode } from "@/data/services/api";

const markPaymentUnlockPendingAction = sendMachineSetup.assign(() => ({
  paymentUnlockPending: true,
  unlockHistoryError: null,
}));

const showUnlockHistoryPromptAction = sendMachineSetup.assign(
  ({ context }) => ({
    paymentUnlockPending: false,
    unlockHistoryPromptVisible: true,
    lockedHistoryCount: countLockedHistoryMessages(context.messages),
    unlockHistoryError: null,
    unlockHistoryPaymentGuidance: null,
  }),
);

const dismissUnlockHistoryPromptAction = sendMachineSetup.assign(() => ({
  paymentUnlockPending: false,
  unlockHistoryPromptVisible: false,
  unlockHistoryError: null,
  unlockHistoryPaymentGuidance: null,
}));

const markUnlockHistoryStartedAction = sendMachineSetup.assign(() => ({
  paymentUnlockPending: false,
  unlockHistoryPromptVisible: false,
  unlockHistoryError: null,
  unlockHistoryPaymentGuidance: null,
}));

const markUnlockHistoryFailedAction = sendMachineSetup.assign(() => ({
  unlockHistoryPromptVisible: true,
  unlockHistoryError: "Failed to unlock messages. Please try again.",
}));

const markUnlockMessageStartedAction = sendMachineSetup.assign(
  ({ event }) => {
    if (event.type !== "ChatUnlockMessageRequested") return {};
    return {
      unlockingMessage: {
        displayMessageId: event.displayMessageId,
        ...(event.remoteMessageId
          ? { messageId: event.remoteMessageId }
          : {}),
        kind: event.kind,
        ...(event.lockType ? { lockType: event.lockType } : {}),
        ...(event.clientLockId
          ? { clientLockId: event.clientLockId }
          : {}),
      },
      unlockMessageError: null,
      unlockPaywallRequest: null,
    };
  },
);

const applyUnlockMessageSucceededAction = sendMachineSetup.assign(
  ({ context, event }) => {
    const { output } = event as unknown as DoneActorEvent<UnlockMessageOutput>;
    return {
      messages: applySingleUnlockOutput(context.messages, output),
      promotion: applyPromotionUnlockOutput(context.promotion, output),
      unlockingMessage: null,
      unlockMessageError: null,
      unlockPaywallRequest: null,
      creditBalance: output.response.creditBalance,
      creditsCharged: output.response.creditsCharged,
      requiredCredits: 0,
      shortfallCredits: 0,
      paymentGuidance: null,
    };
  },
);

const requestUnlockPaymentFromOutputAction = sendMachineSetup.assign(
  ({ context, event }) => {
    const { output } = event as unknown as DoneActorEvent<UnlockMessageOutput>;
    const shouldOpenPaywall = output.response.reason !== "not_found";
    return {
      promotion: applyPromotionUnlockOutput(context.promotion, output),
      unlockingMessage: null,
      unlockMessageError: output.response.reason,
      unlockPaywallRequest: shouldOpenPaywall
        ? {
            displayMessageId: output.displayMessageId,
            ...(output.response.messageId || output.request.messageId
              ? {
                  messageId:
                    output.response.messageId || output.request.messageId,
                }
              : {}),
            kind: context.unlockingMessage?.kind ?? "private",
            ...(output.request.lockType
              ? { lockType: output.request.lockType }
              : {}),
            ...(output.request.clientLockId
              ? { clientLockId: output.request.clientLockId }
              : {}),
            ...(context.promotion?.message.displayId ===
              output.displayMessageId
              ? { promotion: context.promotion.session }
              : {}),
            reason: output.response.reason,
            creditBalance: output.response.creditBalance,
            requiredCredits: output.response.requiredCredits,
            shortfallCredits: output.response.shortfallCredits,
            ...(output.response.paymentGuidance
              ? { paymentGuidance: output.response.paymentGuidance }
              : {}),
          }
        : null,
      creditBalance: output.response.creditBalance,
      requiredCredits: output.response.requiredCredits,
      shortfallCredits: output.response.shortfallCredits,
      paymentGuidance: output.response.paymentGuidance,
    };
  },
);

const requestUnlockPaymentFromErrorAction = sendMachineSetup.assign(
  ({ context }) => {
    const request = context.unlockingMessage;
    return {
      unlockingMessage: null,
      unlockMessageError: "unlock_failed",
      unlockPaywallRequest: request
        ? {
            ...request,
            ...(context.promotion?.message.displayId ===
              request.displayMessageId
              ? { promotion: context.promotion.session }
              : {}),
            reason: "unlock_failed",
            creditBalance: context.creditBalance,
            requiredCredits: context.requiredCredits,
            shortfallCredits: context.shortfallCredits,
          }
        : null,
    };
  },
);

const clearUnlockPaywallRequestAction = sendMachineSetup.assign(() => ({
  unlockPaywallRequest: null,
}));

const handleCharacterMismatchAction = sendMachineSetup.assign(() => ({
  unlockingMessage: null,
  unlockMessageError: "character_mismatch",
  unlockPaywallRequest: null,
  characterErrorCode: "CHARACTER_MISMATCH" as const,
}));

export const unlockMachineSetup = sendMachineSetup.extend({
  actions: {
    markPaymentUnlockPending: markPaymentUnlockPendingAction,
    showUnlockHistoryPrompt: showUnlockHistoryPromptAction,
    dismissUnlockHistoryPrompt: dismissUnlockHistoryPromptAction,
    markUnlockHistoryStarted: markUnlockHistoryStartedAction,
    markUnlockHistoryFailed: markUnlockHistoryFailedAction,
    markUnlockMessageStarted: markUnlockMessageStartedAction,
    applyUnlockMessageSucceeded: applyUnlockMessageSucceededAction,
    requestUnlockPaymentFromOutput: requestUnlockPaymentFromOutputAction,
    requestUnlockPaymentFromError: requestUnlockPaymentFromErrorAction,
    clearUnlockPaywallRequest: clearUnlockPaywallRequestAction,
    handleCharacterMismatch: handleCharacterMismatchAction,
  },
});

const unlockHistoryDoneActionSetup =
  createChatActorActionSetup<DoneActorEvent<UnlockHistoryOutput>>();

const applyUnlockHistoryOutputAction =
  unlockHistoryDoneActionSetup.assign(({ event }) => {
    const lockedHistoryCount = countLockedHistoryMessages(
      event.output.messages,
    );
    const insufficientBalance =
      !event.output.unlocked &&
      event.output.reason === "insufficient_balance";
    return {
      messages: event.output.messages,
      historyLoaded: true,
      paymentUnlockPending: false,
      unlockHistoryPromptVisible: insufficientBalance,
      lockedHistoryCount,
      unlockHistoryError: insufficientBalance
        ? `Insufficient credits. You still need ${event.output.shortfallCredits} credits.`
        : null,
      unlockHistoryPaymentGuidance: insufficientBalance
        ? event.output.paymentGuidance ?? null
        : null,
    };
  });

export const unlockingHistoryState = unlockMachineSetup.createStateConfig({
  invoke: {
    id: "unlockHistory",
    src: "unlockHistory",
    input: ({ context }) => ({
      characterId: context.characterId,
      emptyChatGreeting: context.emptyChatGreeting,
    }),
    onDone: {
      target: "ready",
      actions: applyUnlockHistoryOutputAction,
    },
    onError: {
      target: "ready",
      actions: "markUnlockHistoryFailed",
    },
  },
});

export const unlockingMessageState = unlockMachineSetup.createStateConfig({
  invoke: {
    id: "unlockMessage",
    src: "unlockMessage",
    input: ({ context }) => ({
      characterId: context.characterId,
      ...(context.unlockingMessage ?? {
        displayMessageId: "",
        kind: "private" as const,
      }),
    }),
    onDone: [
      {
        guard: ({ event }) => event.output.response.unlocked,
        target: "ready",
        actions: "applyUnlockMessageSucceeded",
      },
      {
        target: "ready",
        actions: "requestUnlockPaymentFromOutput",
      },
    ],
    onError: [
      {
        guard: ({ event }) =>
          getCharacterErrorCode(event.error) === "CHARACTER_MISMATCH",
        target: "ready",
        actions: "handleCharacterMismatch",
      },
      {
        target: "ready",
        actions: "requestUnlockPaymentFromError",
      },
    ],
  },
});