cozsweet-frontend-nextjs/src/stores/chat/chat-send-state.ts

import type { UiMessage } from "@/data/dto/chat";
import { ChatSendResponse } from "@/data/dto/chat/chat_send_response";

import type { ChatState, ChatUpgradeReason } from "./chat-state";

export type HttpSendOutput = {
  response: ChatSendResponse;
  reply: UiMessage | null;
};

export function beginPendingReply(context: ChatState): Pick<
  ChatState,
  "isReplyingAI" | "pendingReplyCount"
> {
  return {
    pendingReplyCount: context.pendingReplyCount + 1,
    isReplyingAI: true,
  };
}

export function finishPendingReply(context: ChatState): Pick<
  ChatState,
  "isReplyingAI" | "pendingReplyCount"
> {
  const pendingReplyCount = Math.max(0, context.pendingReplyCount - 1);
  return {
    pendingReplyCount,
    isReplyingAI: pendingReplyCount > 0,
  };
}

export function applyHttpSendOutput(
  context: ChatState,
  output: HttpSendOutput,
): Partial<ChatState> {
  const { response, reply } = output;

  const lockDetail = response.lockDetail;
  const isWeeklyLimitBlocked =
    lockDetail.locked &&
    lockDetail.showUpgrade &&
    lockDetail.reason === "weekly_limit";
  const isInsufficientCredits =
    response.canSendMessage === false &&
    response.cannotSendReason === "insufficient_credits";
  const upgradeReason: ChatUpgradeReason | null = isWeeklyLimitBlocked
    ? "weekly_limit"
    : isInsufficientCredits
      ? "insufficient_credits"
      : null;
  const sendCapabilityState = getSendCapabilityState(response);

  if (upgradeReason) {
    const lastMessage = context.messages[context.messages.length - 1];
    const messages =
      !reply && lastMessage && !lastMessage.isFromAI
        ? context.messages.slice(0, -1)
        : context.messages;

    return {
      messages: reply ? [...messages, reply] : messages,
      ...finishPendingReply(context),
      upgradePromptVisible: true,
      upgradeReason,
      ...sendCapabilityState,
    };
  }

  if (!reply) {
    return {
      ...finishPendingReply(context),
      ...sendCapabilityState,
    };
  }

  if (lockDetail.showUpgrade && lockDetail.reason === "private_message") {
    return {
      messages: [...context.messages, reply],
      ...finishPendingReply(context),
      upgradePromptVisible: false,
      upgradeReason: null,
      ...sendCapabilityState,
    };
  }

  return {
    messages: [...context.messages, reply],
    ...finishPendingReply(context),
    upgradePromptVisible: false,
    upgradeReason: null,
    ...sendCapabilityState,
  };
}

function getSendCapabilityState(
  response: ChatSendResponse,
): Pick<
  ChatState,
  | "canSendMessage"
  | "cannotSendReason"
  | "creditBalance"
  | "creditsCharged"
  | "requiredCredits"
  | "shortfallCredits"
> {
  return {
    canSendMessage: response.canSendMessage,
    cannotSendReason: response.cannotSendReason,
    creditBalance: response.creditBalance,
    creditsCharged: response.creditsCharged,
    requiredCredits: response.requiredCredits,
    shortfallCredits: response.shortfallCredits,
  };
}