storytellurx-v1/app/components/chat/chat-message.tsx

"use client";

import { type VariantProps } from "class-variance-authority";
import { toConstantCase } from "@std/text/unstable-to-constant-case";
import cn from "@bureaudouble/outils/cn.ts";
import { markdownParse } from "../../utils/markdown-parser.ts";
import { LLMRequestViewer } from "./llm-request-viewer.tsx";
import InfoIcon from "@bureaudouble/icons/lucide/Info";
import { Button } from "../ui/button.tsx";
import { ImageDisplay } from "../chat/chat-image.tsx";
import { useTurn } from "../../game-engine/turn-provider.tsx";
import {
  gameStateToVariantState,
  personaMessageContainerVariants,
  personaMessageVariants,
  personaRationaleVariants,
  personaStateChangeVariants,
} from "../../game-engine/persona-styles.ts";
import { useNarrative } from "../../game-engine/narrative-provider.tsx";
import { useClientLocale } from "../../game-engine/locale-provider.tsx";

export const Rationale = (
  { content, title }: { content?: string | null; title?: string },
) => {
  const { t } = useClientLocale();
  const { gameState } = useNarrative();
  const { nextAssistantMessage } = useTurn();
  const variantState = gameStateToVariantState(gameState);

  const i18nKey = title?.toLowerCase().replace(
    /_(\w)/g,
    (_, letter) => letter.toUpperCase(),
  );
  const translatedTitle = i18nKey && t(i18nKey as any)
    ? t(i18nKey as any)
    : title?.split("_").join(" ");

  return content && (
    <div
      tabIndex={0}
      data-is-last-turn={!nextAssistantMessage || undefined}
      className={cn("text-left", personaRationaleVariants(variantState))}
    >
      <div data-container>
        <div data-title>{translatedTitle}</div>
        <div
          data-contents
          className="contents"
          dangerouslySetInnerHTML={{ __html: content }}
        />
      </div>
    </div>
  );
};

interface MessageProps
  extends
    Omit<React.HTMLAttributes<HTMLDivElement>, "role">,
    VariantProps<typeof personaMessageVariants> {
  content: string;
  role?: "system" | "user" | "assistant" | null;
}

const Message: React.FC<MessageProps> = ({ role, className, content }) => {
  const { message } = useTurn();
  const { gameState } = useNarrative();
  const variantState = gameStateToVariantState(gameState);

  return (
    <div
      data-turn-role={message.role}
      data-turn-id={message.id}
      className={cn(
        personaMessageVariants({ ...variantState, role }),
        className,
      )}
      dangerouslySetInnerHTML={{ __html: content }}
    />
  );
};

const StateChangeIndicator = () => {
  const { turnData, previousTurnData } = useTurn();
  const currentGameState = turnData?.fullGameState;
  const previousGameState = previousTurnData?.fullGameState;
  if (!currentGameState || !previousGameState) return null;

  const stateTransitionChanges = Object.entries(currentGameState)
    .filter(([_k, v]) => v !== undefined && v !== null)
    .map(([key, value]) => {
      const previous = previousGameState[key as keyof typeof previousGameState];
      return {
        key,
        current: Array.isArray(value) && value.length === 0 ? undefined : value,
        previous: Array.isArray(previous) && previous.length === 0
          ? undefined
          : previous,
      };
    })
    .filter((obj) =>
      JSON.stringify(obj.current) !== JSON.stringify(obj.previous)
    );

  return (
    <div
      className={personaStateChangeVariants(
        gameStateToVariantState(turnData.fullGameState),
      )}
    >
      {stateTransitionChanges.map((obj) => (
        <div key={obj.key}>
          {`// ${toConstantCase(obj.key)} :: ${
            ["string", "number"].includes(typeof obj.current)
              ? obj.current
              : JSON.stringify(obj.current, null, 1)
          }`}
        </div>
      ))}
    </div>
  );
};

const AssistantMessageDisplay: React.FC<{ className?: string }> = (
  { className },
) => {
  const { t } = useClientLocale();
  const { message, previousAssistantMessage, turnData, previousMessages } =
    useTurn();
  if (!turnData) return null;

  const { aiMessagesArrayToUser, fullGameState, aiStateChangeRationale } =
    turnData;
  const { personaId } = fullGameState;

  const displayContent =
    (Array.isArray(aiMessagesArrayToUser)
      ? aiMessagesArrayToUser
      : typeof aiMessagesArrayToUser === "string"
      ? [aiMessagesArrayToUser]
      : [aiMessagesArrayToUser]).map((
        v,
      ) => (typeof ((v as any)?.text ?? null) === "string"
        ? (v as any).text
        : v)
      ).filter((v) => (typeof v === "string" ? v : JSON.stringify(v)));
  const isAISystemStyleMessage =
    (!!aiStateChangeRationale && displayContent?.length === 0) ||
    personaId === "storytellurx_gateway" || personaId === "storytellurx_pause";

  const rationaleItems = [
    { content: turnData.aiStateChangeRationale, label: "STATE_RATIONALE" },
    { content: turnData.aiPlayRationale, label: "PLAY_RATIONALE" },
    { content: turnData.aiContextRationale, label: "CONTEXT_RATIONALE" },
    { content: turnData.aiCTARationale, label: "CTA_RATIONALE" },
  ];

  const role = isAISystemStyleMessage ? "system" : "assistant";

  return (
    <div className={cn("relative", className)}>
      {previousAssistantMessage && <StateChangeIndicator />}
      {rationaleItems.map((item, index) => (
        <Rationale key={index} title={item.label} content={item.content} />
      ))}
      <ImageDisplay message={message} turnData={turnData} />
      {displayContent.map((content) => (
        <Message
          key={content}
          role={role}
          content={`${
            personaId === "storytellurx_pause" &&
              fullGameState?.personaId !== "storytellurx_pause"
              ? "[REDACTED]"
              : markdownParse(content)
          }`}
        />
      ))}
      {previousMessages && (
        <div className="self-start absolute -bottom-1 right-0 not-hover:opacity-90">
          <LLMRequestViewer messages={[...previousMessages, message]}>
            <Button
              variant="ghost"
              size="sm"
              title={t("llmRequestPayload")}
            >
              <InfoIcon className="size-3" />
            </Button>
          </LLMRequestViewer>
        </div>
      )}
    </div>
  );
};

const RewrittenUserPrompt = () => {
  const { t } = useClientLocale();
  const { message, nextTurnData: nextTurn } = useTurn();

  return nextTurn?.mindControlledRewritenUserPrompt &&
      nextTurn?.mindControlledRewritenUserPrompt !==
        message.parts.filter((v) => v.type === "text").at(0)?.text
    ? (
      <div className="p-1.5 mr-auto ml-1 border border-dashed border-red-500 max-w-[90cqw] self-start text-xs">
        <div className="text-red-400 font-bold">
          {t("aiIntervention")}
        </div>
        <div className="text-muted-foreground mt-0.5">
          <span className="font-semibold">{t("reason")}</span>{" "}
          {nextTurn?.aiRewriteRationale ?? t("notAvailable")}
        </div>
      </div>
    )
    : nextTurn?.aiRewriteRationale
    ? (
      <Rationale
        title={t("rewriteDecision")}
        content={nextTurn?.aiRewriteRationale}
      />
    )
    : null;
};

const UserMessageDisplay: React.FC<{ className?: string }> = (
  { className },
) => {
  const { message, nextTurnData: nextTurn } = useTurn();

  return message.parts
    .filter((part) => part.type === "text")
    .map((part, i) => (
      <div className={className} key={`${part.text}-${i}`}>
        <Message
          role="user"
          content={nextTurn?.mindControlledRewritenUserPrompt ?? part.text}
        />
        <RewrittenUserPrompt />
      </div>
    ));
};

export default function ChatMessageDisplay() {
  const { gameState } = useNarrative();
  const { message } = useTurn();
  const Component = (() => {
    switch (message.role) {
      case "user":
        return UserMessageDisplay;
      case "assistant":
        return AssistantMessageDisplay;
      default:
        return null;
    }
  })();
  return Component && (
    <Component
      className={cn(
        personaMessageContainerVariants({
          ...gameStateToVariantState(gameState),
          role: message.role,
        }),
      )}
    />
  );
}