Jyotisha/frontend/src/components/thinking-step-tree.tsx

"use client";

import { useEffect, useId, useState } from "react";
import dynamic from "next/dynamic";
import { Check } from "lucide-react";
import type { OrbState } from "thinking-orbs";

import { prefetchOnIdle } from "@/components/chat-chunk-prefetch";
import { activityElapsedLabel } from "@/lib/chat-message-view";
import {
  visibleThinkingSteps,
  type PublicThinkingStep,
  type ThinkingStepStatus,
} from "@/lib/consultation-thinking-plan";

const importThinkingOrb = () => import("thinking-orbs");

const ThinkingOrb = dynamic(async () => (await importThinkingOrb()).ThinkingOrb, {
  loading: () => (
    <span aria-hidden="true" style={{ display: "block", flex: "0 0 auto", height: 20, width: 20 }} />
  ),
  ssr: false,
});

prefetchOnIdle(importThinkingOrb);

function ActivityElapsed({ startedAt }: Readonly<{ startedAt: number }>) {
  const [now, setNow] = useState(() => Date.now());
  useEffect(() => {
    const timer = window.setInterval(() => setNow(Date.now()), 1000);
    return () => window.clearInterval(timer);
  }, []);
  const label = activityElapsedLabel(startedAt, now);
  if (!label) return null;
  return <span className="agent-activity-status__elapsed" aria-hidden="true">{label}</span>;
}

function StepMarker({
  status,
  liveState,
}: Readonly<{
  status: ThinkingStepStatus | "live";
  liveState?: OrbState;
}>) {
  if (status === "live" || status === "active") {
    return (
      <span className="agent-thinking-marker is-live-marker">
        <ThinkingOrb aria-hidden="true" state={liveState ?? "working"} size={20} />
      </span>
    );
  }
  if (status === "done") {
    return (
      <span className="agent-thinking-marker" aria-hidden="true">
        <Check />
      </span>
    );
  }
  return <span className="agent-thinking-marker is-pending" aria-hidden="true" />;
}

type ThinkingStepGroup = Readonly<{
  id: string;
  intent?: string;
  steps: readonly PublicThinkingStep[];
  live?: boolean;
}>;

function StepList({
  steps,
  revealAll,
  live,
  liveLabel,
  liveState,
  startedAt,
}: Readonly<{
  steps: readonly PublicThinkingStep[];
  revealAll: boolean;
  live: boolean;
  liveLabel?: string;
  liveState?: OrbState;
  startedAt?: number;
}>) {
  const visible = revealAll ? { visible: steps, hiddenCount: 0 } : visibleThinkingSteps(steps);
  const hiddenDone = visible.hiddenCount > 0
    && steps.slice(visible.visible.length).every((step) => step.status === "done");
  return (
    <ol className="agent-thinking-timeline">
      {visible.visible.map((step) => (
        <li
          className={`agent-thinking-step is-${step.status}`}
          key={step.id}
        >
          <StepMarker status={step.status} liveState={liveState} />
          <span>{step.label}</span>
        </li>
      ))}
      {visible.hiddenCount > 0 ? (
        <li className={`agent-thinking-step is-more${hiddenDone ? " is-done" : ""}`}>
          <StepMarker status={hiddenDone ? "done" : "pending"} />
          <span>还有 {visible.hiddenCount} 项</span>
        </li>
      ) : null}
      {live && liveLabel ? (
        <li className="agent-thinking-step is-live">
          <StepMarker status="live" liveState={liveState} />
          <span className="agent-activity-status__live" role="status">
            <span key={liveLabel} className="agent-activity-status__text">{liveLabel}</span>
            {startedAt ? <ActivityElapsed key={startedAt} startedAt={startedAt} /> : null}
          </span>
        </li>
      ) : null}
    </ol>
  );
}

export function ThinkingStepTree({
  caption,
  intent,
  reasoning,
  steps = [],
  groups,
  hiddenCount,
  revealAll = false,
  live = false,
  liveLabel,
  liveState = "working",
  startedAt,
  defaultOpen = false,
}: Readonly<{
  caption?: string;
  intent?: string;
  reasoning?: string;
  steps?: readonly PublicThinkingStep[];
  groups?: readonly ThinkingStepGroup[];
  hiddenCount?: number;
  revealAll?: boolean;
  live?: boolean;
  liveLabel?: string;
  liveState?: OrbState;
  startedAt?: number;
  defaultOpen?: boolean;
}>) {
  const [userOpen, setUserOpen] = useState<boolean | null>(null);
  const treeId = useId();
  const groupItems = groups ?? [{
    id: "default",
    intent,
    steps,
    live,
  }];
  const showAll = revealAll || hiddenCount === 0;
  const stagedCount = groupItems.filter((group) => Boolean(group.intent)).length;
  const body = (
    <>
      {groupItems.map((group, index) => {
        const stageId = `${treeId}-${group.id}-stage`;
        return (
          <section
            className="consultation-step-tree__group"
            key={group.id}
            aria-labelledby={group.intent ? stageId : undefined}
          >
            {group.intent ? (
              <h3 className="consultation-step-tree__stage" id={stageId}>
                {stagedCount > 1 ? (
                  <span className="consultation-step-tree__stage-index" aria-hidden="true">
                    {index + 1}
                  </span>
                ) : (
                  <span className="consultation-step-tree__stage-mark" aria-hidden="true" />
                )}
                <span className="consultation-step-tree__stage-title">
                  {stagedCount > 1 ? <span className="sr-only">阶段 {index + 1}：</span> : null}
                  {group.intent}
                </span>
              </h3>
            ) : null}
            <StepList
              steps={group.steps}
              revealAll={showAll}
              live={Boolean(group.live)}
              liveLabel={group.live ? liveLabel : undefined}
              liveState={liveState}
              startedAt={group.live ? startedAt : undefined}
            />
          </section>
        );
      })}
    </>
  );

  if (!caption) return <div className="consultation-step-tree">{body}</div>;

  const open = userOpen ?? defaultOpen;
  return (
    <details
      className="message-thinking consultation-step-tree"
      open={open}
      onToggle={(event) => {
        setUserOpen((event.currentTarget as HTMLDetailsElement).open);
      }}
    >
      <summary>{caption}</summary>
      <div className="message-thinking-body">
        {reasoning?.trim() ? (
          <div className="consultation-step-tree__reasoning">{reasoning}</div>
        ) : null}
        {body}
      </div>
    </details>
  );
}