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

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 isInsufficientCredits = response.canSendMessage === false;
  const upgradeReason: ChatUpgradeReason | null = 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,
    };
  }

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

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