Jyotisha/frontend/src/components/chat-message-content.tsx

"use client";

import { memo, useEffect, useState, type ReactNode } from "react";

import { prefetchOnIdle } from "@/components/chat-chunk-prefetch";
import { plainParagraphs } from "@/components/chat-message-paragraphs";
import { TechniqueAuditDisclosure } from "@/components/technique-audit-disclosure";
import { splitSpokenAndReport } from "@/lib/chat-answer-split";
import { splitStableMarkdown } from "@/lib/chat-markdown-split";
import type { TechniqueAuditRow } from "@/lib/consultation-agent-events";
import {
  resolveTechniqueAuditRows,
  splitSpokenAnswerAndTechniqueAudit,
} from "@/lib/consultation-technique-audit";

type MarkdownRenderer = (text: string) => ReactNode;

let markdownRenderer: MarkdownRenderer | undefined;
let markdownRendererRequest: Promise<MarkdownRenderer> | undefined;

function loadMarkdownRenderer() {
  markdownRendererRequest ??= import("@/components/chat-markdown-view").then((module) => {
    markdownRenderer = module.renderChatMarkdown;
    return module.renderChatMarkdown;
  });
  return markdownRendererRequest;
}

prefetchOnIdle(loadMarkdownRenderer);

function useMarkdownRenderer() {
  const [renderer, setRenderer] = useState<MarkdownRenderer | undefined>(() => markdownRenderer);

  useEffect(() => {
    if (renderer) return;
    let active = true;
    void loadMarkdownRenderer().then((loaded) => {
      if (active) setRenderer(() => loaded);
    });
    return () => { active = false; };
  }, [renderer]);

  return renderer;
}

function renderProse(text: string, renderMarkdown: MarkdownRenderer | undefined): ReactNode {
  if (!text) return null;
  return renderMarkdown
    ? renderMarkdown(text)
    : (plainParagraphs(text) ?? []).map((paragraph, index) => (
      <p key={index}>{paragraph}</p>
    ));
}

const settledMarkdownNodes = new WeakMap<MarkdownRenderer, Map<string, ReactNode>>();

function rememberedProse(text: string, renderMarkdown: MarkdownRenderer | undefined): ReactNode {
  if (!text) return null;
  if (!renderMarkdown) return renderProse(text, undefined);
  let byText = settledMarkdownNodes.get(renderMarkdown);
  if (!byText) {
    byText = new Map();
    settledMarkdownNodes.set(renderMarkdown, byText);
  }
  const cached = byText.get(text);
  if (cached !== undefined) return cached;
  const node = renderMarkdown(text);
  byText.set(text, node);
  return node;
}

/**
 * Settled prose and the completed prefix of a streaming answer. `memo` plus a
 * content-keyed parse cache keep React from calling the markdown parser again
 * while `text` is unchanged, so a parent re-render (or a frame that only grew
 * the tail) does not parse the whole answer.
 */
export const StableMarkdownPrefix = memo(function StableMarkdownPrefix({
  text,
  renderMarkdown,
}: Readonly<{
  text: string;
  renderMarkdown: MarkdownRenderer | undefined;
}>) {
  return <>{rememberedProse(text, renderMarkdown)}</>;
});

export function StreamingMarkdown({
  text,
  renderMarkdown,
}: Readonly<{
  text: string;
  renderMarkdown: MarkdownRenderer | undefined;
}>) {
  const split = splitStableMarkdown(text);
  return (
    <>
      {split.stable ? <StableMarkdownPrefix text={split.stable} renderMarkdown={renderMarkdown} /> : null}
      {renderProse(split.tail, renderMarkdown)}
    </>
  );
}

export function ChatMessageContent({
  text,
  auditRows,
  streaming = false,
}: {
  text: string;
  auditRows?: readonly TechniqueAuditRow[];
  streaming?: boolean;
}) {
  const renderMarkdown = useMarkdownRenderer();
  const split = splitSpokenAnswerAndTechniqueAudit(text);
  const rows = resolveTechniqueAuditRows(split, auditRows);
  const layers = splitSpokenAndReport(split.spoken);
  const spoken = layers.spoken;
  const report = layers.report;
  const showFoldedAudit = rows.length > 0 && split.rows.length === 0;
  const foldedAudit = showFoldedAudit ? <TechniqueAuditDisclosure rows={rows} /> : null;
  const reportBody = !streaming && report
    ? (
      <>
        <div className="message-markdown">
          <StableMarkdownPrefix text={report} renderMarkdown={renderMarkdown} />
        </div>
        {foldedAudit}
      </>
    )
    : null;

  return (
    <div className="message-answer">
      {spoken ? (
        <div className="message-markdown">
          {streaming
            ? <StreamingMarkdown text={spoken} renderMarkdown={renderMarkdown} />
            : <StableMarkdownPrefix text={spoken} renderMarkdown={renderMarkdown} />}
        </div>
      ) : null}
      {/* The body used to sit inside a collapsed 完整分析 control: everything
          from the first H2 down was one click away, so the part of the answer
          that carries the reasoning was invisible by default. The answer is the
          product — it renders inline. The technique audit keeps its own
          disclosure, because that one is evidence for checking, not the reply. */}
      {report ? <div className="answer-report-body">{reportBody}</div> : foldedAudit}
    </div>
  );
}