socaichat/frontend/components/gemini/GeminiChat.tsx

"use client";

import { useState, useRef, useEffect, useCallback } from "react";
import { GeminiSidebar } from "./GeminiSidebar";
import { GeminiTopbar } from "./GeminiTopbar";
import { GeminiWelcome } from "./GeminiWelcome";
import { GeminiMessage, Message } from "./GeminiMessage";
import { GeminiInput } from "./GeminiInput";
import { GeminiTypingIndicator } from "./GeminiTypingIndicator";
import { ExtensionsPanel, TicketSummary } from "./ExtensionsPanel";
import { Ticket, ShoppingCart, Cloud } from "lucide-react";
import { cn } from "@/lib/utils";
import {
  fetchConversations,
  fetchConversation,
  fetchTicketSummary,
  streamChat,
  deleteConversation,
  type TicketSummaryData,
  type AttachmentData,
  type TraceItem,
  type WorkspaceSession,
  type ActivityNode,
  type WorkspaceCard,
} from "@/lib/api";
import { AgentWorkspace } from "@/components/workspace/AgentWorkspace";

// ── Types ────────────────────────────────────────────────────────────────────
interface Conversation {
  id: string;
  title: string;
  messages: Message[];
}

interface Extension {
  id: string;
  name: string;
  description: string;
  icon: React.ElementType;
  connected: boolean;
  apiKey: string;
}

const INITIAL_EXTENSIONS: Extension[] = [
  {
    id: "ticket",
    name: "工单系统",
    description: "连接工单系统，查看和管理技术支持工单",
    icon: Ticket,
    connected: false,
    apiKey: "",
  },
  {
    id: "sales",
    name: "销售系统",
    description: "连接销售 CRM，获取客户和订单数据",
    icon: ShoppingCart,
    connected: false,
    apiKey: "",
  },
  {
    id: "cloud",
    name: "云管系统",
    description: "连接云管平台，监控资源使用情况",
    icon: Cloud,
    connected: false,
    apiKey: "",
  },
];

// ── Trace helpers ─────────────────────────────────────────────────────────────

function appendTraceItem(
  convId: string,
  msgId: string,
  item: TraceItem,
  setConversations: React.Dispatch<React.SetStateAction<Conversation[]>>
) {
  setConversations((prev) =>
    prev.map((c) => {
      if (c.id !== convId) return c;
      return {
        ...c,
        messages: c.messages.map((m) =>
          m.id === msgId
            ? { ...m, traceItems: [...(m.traceItems ?? []), item] }
            : m
        ),
      };
    })
  );
}

function updateTraceItem(
  convId: string,
  msgId: string,
  callId: string | undefined,
  tool: string,
  updates: Partial<TraceItem>,
  setConversations: React.Dispatch<React.SetStateAction<Conversation[]>>
) {
  setConversations((prev) =>
    prev.map((c) => {
      if (c.id !== convId) return c;
      return {
        ...c,
        messages: c.messages.map((m) => {
          if (m.id !== msgId) return m;
          const items = [...(m.traceItems ?? [])];
          for (let i = items.length - 1; i >= 0; i--) {
            // Prefer precise callId match
            if (callId && items[i].callId === callId) {
              items[i] = { ...items[i], ...updates };
              break;
            }
            // Fallback: last running entry for same tool name
            if (!callId && items[i].tool === tool && items[i].itemStatus === "running") {
              items[i] = { ...items[i], ...updates };
              break;
            }
          }
          return { ...m, traceItems: items };
        }),
      };
    })
  );
}

// ── Workspace helpers ─────────────────────────────────────────────────────────

function initWorkspaceSession(
  aiMsgId: string,
  convId: string,
  title: string,
  setWS: React.Dispatch<React.SetStateAction<Map<string, WorkspaceSession>>>
) {
  setWS(prev => {
    if (prev.has(aiMsgId)) return prev;
    const next = new Map(prev);
    next.set(aiMsgId, {
      id: aiMsgId, conversationId: convId, messageId: aiMsgId,
      title: title.slice(0, 40), status: "running", stageLabel: "正在分析...",
      timeline: [], cards: [], startedAt: Date.now(),
    });
    return next;
  });
}

function addWsActivityNode(
  aiMsgId: string,
  node: ActivityNode,
  setWS: React.Dispatch<React.SetStateAction<Map<string, WorkspaceSession>>>
) {
  setWS(prev => {
    const s = prev.get(aiMsgId);
    if (!s) return prev;
    const next = new Map(prev);
    next.set(aiMsgId, { ...s, timeline: [...s.timeline, node] });
    return next;
  });
}

function updateWsActivityNode(
  aiMsgId: string,
  callId: string | undefined,
  tool: string,
  updates: Partial<ActivityNode>,
  setWS: React.Dispatch<React.SetStateAction<Map<string, WorkspaceSession>>>
) {
  setWS(prev => {
    const s = prev.get(aiMsgId);
    if (!s) return prev;
    const timeline = s.timeline.map(n => {
      if (callId && n.callId === callId) return { ...n, ...updates };
      if (!callId && n.tool === tool && n.nodeStatus === "running") return { ...n, ...updates };
      return n;
    });
    const next = new Map(prev);
    next.set(aiMsgId, { ...s, timeline });
    return next;
  });
}

function addWsCard(
  aiMsgId: string,
  card: WorkspaceCard,
  setWS: React.Dispatch<React.SetStateAction<Map<string, WorkspaceSession>>>
) {
  setWS(prev => {
    const s = prev.get(aiMsgId);
    if (!s) return prev;
    const next = new Map(prev);
    next.set(aiMsgId, { ...s, cards: [...s.cards, card] });
    return next;
  });
}

function mergeWsCard(
  aiMsgId: string,
  cardId: string,
  props: Record<string, unknown>,
  setWS: React.Dispatch<React.SetStateAction<Map<string, WorkspaceSession>>>
) {
  setWS(prev => {
    const s = prev.get(aiMsgId);
    if (!s) return prev;
    const cards = s.cards.map(c => c.id === cardId ? { ...c, props: { ...c.props, ...props } } : c);
    const next = new Map(prev);
    next.set(aiMsgId, { ...s, cards });
    return next;
  });
}

function completeWsSession(
  aiMsgId: string,
  setWS: React.Dispatch<React.SetStateAction<Map<string, WorkspaceSession>>>
) {
  setWS(prev => {
    const s = prev.get(aiMsgId);
    if (!s) return prev;
    const timeline = s.timeline.map(n => n.nodeStatus === "running" ? { ...n, nodeStatus: "success" as const } : n);
    const next = new Map(prev);
    next.set(aiMsgId, { ...s, status: "completed", stageLabel: "已完成", timeline, finishedAt: Date.now() });
    return next;
  });
}

// ── Main Component ────────────────────────────────────────────────────────────
export function GeminiChat() {
  const [sidebarOpen, setSidebarOpen] = useState(true);
  const [conversations, setConversations] = useState<Conversation[]>([]);
  const [activeConvId, setActiveConvId] = useState<string | null>(null);
  const [inputValue, setInputValue] = useState("");
  const [isLoading, setIsLoading] = useState(false);
  const [extensionsPanelOpen, setExtensionsPanelOpen] = useState(false);
  const [extensions, setExtensions] = useState<Extension[]>(INITIAL_EXTENSIONS);
  const [ticketSummary, setTicketSummary] = useState<TicketSummaryData | null>(null);
  const [activeTools, setActiveTools] = useState<Set<string>>(new Set());
  const [selectedModel, setSelectedModel] = useState<"flash" | "auto" | "pro">("auto");
  const [workspaceSessions, setWorkspaceSessions] = useState<Map<string, WorkspaceSession>>(new Map());
  const [activeWorkspaceId, setActiveWorkspaceId] = useState<string | null>(null);
  const messagesEndRef = useRef<HTMLDivElement>(null);
  const abortRef = useRef<AbortController | null>(null);

  const activeWorkspace = activeWorkspaceId ? (workspaceSessions.get(activeWorkspaceId) ?? null) : null;

  const activeConversation = conversations.find((c) => c.id === activeConvId) ?? null;
  const messages = activeConversation?.messages ?? [];

  // Map UI model to backend model ("auto" -> "flash")
  const apiModel: "flash" | "pro" = selectedModel === "pro" ? "pro" : "flash";

  // 检查工单系统是否已连接
  const ticketSystemConnected = extensions.find((e) => e.id === "ticket")?.connected ?? false;

  // Load conversations on mount
  useEffect(() => {
    fetchConversations()
      .then((list) => {
        setConversations(
          list.map((c) => ({ id: c.id, title: c.title, messages: [] }))
        );
      })
      .catch(() => {
        // silently fail — user sees empty sidebar
      });
  }, []);

  // Load ticket summary when ticket extension connects
  useEffect(() => {
    if (ticketSystemConnected) {
      fetchTicketSummary()
        .then(setTicketSummary)
        .catch(() => setTicketSummary(null));
    } else {
      setTicketSummary(null);
    }
  }, [ticketSystemConnected]);

  // Scroll to bottom on new messages
  useEffect(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
  }, [messages, isLoading]);

  const handleNewChat = useCallback(() => {
    setActiveConvId(null);
    setInputValue("");
  }, []);

  const handleSelectConversation = useCallback(
    (id: string) => {
      setActiveConvId(id);
      setInputValue("");

      // Load messages if not yet loaded and no stream is in progress
      const conv = conversations.find((c) => c.id === id);
      if (conv && conv.messages.length === 0 && !isLoading) {
        fetchConversation(id)
          .then((detail) => {
            const msgs: Message[] = detail.messages.map((m) => ({
              id: m.id,
              role: m.role === "human" ? "user" : "assistant",
              content: m.content,
              timestamp: new Date(m.created_at),
            }));
            // Only write history if the conversation still has no messages
            // (guards against a race where streaming already populated it)
            setConversations((prev) =>
              prev.map((c) =>
                c.id === id && c.messages.length === 0
                  ? { ...c, messages: msgs }
                  : c
              )
            );
          })
          .catch(() => {
            // silently fail
          });
      }
    },
    [conversations, isLoading]
  );

  const handleRenameConversation = useCallback(
    (id: string, title: string) => {
      // Optimistically update locally
      setConversations((prev) =>
        prev.map((c) => (c.id === id ? { ...c, title } : c))
      );
      // PATCH on backend (fire-and-forget)
      fetch(`${process.env.NEXT_PUBLIC_API_URL ?? "https://soc-backend.azurewebsites.net"}/api/conversations/${id}`, {
        method: "PATCH",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ title }),
      }).catch(() => {
        // silently fail — local state already updated
      });
    },
    []
  );

  const handleDeleteConversation = useCallback(
    (id: string) => {
      setConversations((prev) => prev.filter((c) => c.id !== id));
      if (activeConvId === id) {
        setActiveConvId(null);
      }
      deleteConversation(id).catch(() => {
        // silently fail
      });
    },
    [activeConvId]
  );

  const handleOpenExtensions = useCallback(() => {
    setExtensionsPanelOpen(true);
  }, []);

  const handleUpdateExtension = useCallback(
    (id: string, apiKey: string, connected: boolean) => {
      setExtensions((prev) =>
        prev.map((ext) =>
          ext.id === id ? { ...ext, apiKey, connected } : ext
        )
      );
    },
    []
  );

  const handleSend = useCallback(async (attachments?: AttachmentData[]) => {
    const text = inputValue.trim();
    if (!text || isLoading) return;
    setInputValue("");
    setIsLoading(true);

    const userMsg: Message = {
      id: `msg-${Date.now()}-user`,
      role: "user",
      content: text,
      timestamp: new Date(),
      ...(attachments && attachments.length > 0 ? { attachments } : {}),
    };

    let convId = activeConvId;

    if (!convId) {
      // Generate a UUID client-side; backend streamChat auto-creates the conversation
      convId = crypto.randomUUID();
      const title = text.length > 50 ? text.slice(0, 50) + "…" : text;
      const newConv: Conversation = {
        id: convId,
        title,
        messages: [userMsg],
      };
      setConversations((prev) => [newConv, ...prev]);
      setActiveConvId(convId);
    } else {
      setConversations((prev) =>
        prev.map((c) =>
          c.id === convId ? { ...c, messages: [...c.messages, userMsg] } : c
        )
      );
    }

    const aiMsgId = `msg-${Date.now()}-ai`;
    const streamConvId = convId;

    abortRef.current = streamChat(
      text,
      streamConvId,
      Array.from(activeTools),
      apiModel,
      (event) => {
        if (event.type === "token" && event.content) {
          setConversations((prev) =>
            prev.map((c) => {
              if (c.id !== streamConvId) return c;
              const exists = c.messages.some((m) => m.id === aiMsgId);
              if (!exists) {
                // First token: create the assistant message (preserve accumulated traceItems)
                const traceItems = c.messages.find((m) => m.id === aiMsgId)?.traceItems ?? [];
                return {
                  ...c,
                  messages: [
                    ...c.messages,
                    {
                      id: aiMsgId,
                      role: "assistant" as const,
                      content: event.content!,
                      timestamp: new Date(),
                      traceItems,
                    },
                  ],
                };
              }
              // Subsequent tokens: append to existing message
              return {
                ...c,
                messages: c.messages.map((m) =>
                  m.id === aiMsgId
                    ? { ...m, content: m.content + event.content }
                    : m
                ),
              };
            })
          );

        } else if (event.type === "status") {
          // status 事件：确保 assistant 消息占位符已存在
          setConversations((prev) =>
            prev.map((c) => {
              if (c.id !== streamConvId) return c;
              const exists = c.messages.some((m) => m.id === aiMsgId);
              if (exists) return c;
              return {
                ...c,
                messages: [
                  ...c.messages,
                  { id: aiMsgId, role: "assistant" as const, content: "", timestamp: new Date(), traceItems: [] },
                ],
              };
            })
          );
          // workspace: init session + update stageLabel
          initWorkspaceSession(aiMsgId, streamConvId, text, setWorkspaceSessions);
          setActiveWorkspaceId(aiMsgId);
          setWorkspaceSessions(prev => {
            const s = prev.get(aiMsgId);
            if (!s) return prev;
            const next = new Map(prev);
            next.set(aiMsgId, { ...s, stageLabel: event.message ?? "正在处理..." });
            return next;
          });

        } else if (event.type === "tool_start" && event.tool) {
          const item: TraceItem = {
            id: event.call_id ?? `${event.tool}-${event.ts ?? Date.now()}`,
            callId: event.call_id,
            type: "tool_start",
            tool: event.tool,
            title: event.title ?? event.tool,
            inputSummary: event.input_summary,
            itemStatus: "running",
            startTs: event.ts ?? Date.now(),
          };
          // Ensure assistant message placeholder exists
          setConversations((prev) =>
            prev.map((c) => {
              if (c.id !== streamConvId) return c;
              const exists = c.messages.some((m) => m.id === aiMsgId);
              if (!exists) {
                return {
                  ...c,
                  messages: [
                    ...c.messages,
                    { id: aiMsgId, role: "assistant" as const, content: "", timestamp: new Date(), traceItems: [item] },
                  ],
                };
              }
              return {
                ...c,
                messages: c.messages.map((m) =>
                  m.id === aiMsgId ? { ...m, traceItems: [...(m.traceItems ?? []), item] } : m
                ),
              };
            })
          );
          // workspace: init + add timeline node
          initWorkspaceSession(aiMsgId, streamConvId, text, setWorkspaceSessions);
          setActiveWorkspaceId(aiMsgId);
          const wsNode: ActivityNode = {
            id: event.call_id ?? `${event.tool}-${Date.now()}`,
            type: "tool",
            label: event.title ?? event.tool,
            tool: event.tool,
            callId: event.call_id,
            nodeStatus: "running",
            ts: event.ts ?? Date.now(),
          };
          addWsActivityNode(aiMsgId, wsNode, setWorkspaceSessions);
          setWorkspaceSessions(prev => {
            const s = prev.get(aiMsgId);
            if (!s) return prev;
            const next = new Map(prev);
            next.set(aiMsgId, { ...s, stageLabel: `正在调用 ${event.title ?? event.tool}...` });
            return next;
          });

        } else if (event.type === "tool_end" && event.tool) {
          updateTraceItem(streamConvId, aiMsgId, event.call_id, event.tool, {
            type: "tool_end",
            outputSummary: event.output_summary,
            itemStatus: "success",
            durationMs: event.duration_ms,
          }, setConversations);
          updateWsActivityNode(aiMsgId, event.call_id, event.tool, { nodeStatus: "success" }, setWorkspaceSessions);

        } else if (event.type === "tool_error" && event.tool) {
          updateTraceItem(streamConvId, aiMsgId, event.call_id, event.tool, {
            type: "tool_error",
            errorSummary: event.error_summary,
            itemStatus: "error",
            durationMs: event.duration_ms,
          }, setConversations);
          updateWsActivityNode(aiMsgId, event.call_id, event.tool, { nodeStatus: "error" }, setWorkspaceSessions);

        } else if (event.type === "workspace_card" && event.id && event.name) {
          initWorkspaceSession(aiMsgId, streamConvId, text, setWorkspaceSessions);
          if (event.merge) {
            mergeWsCard(aiMsgId, event.id, event.props ?? {}, setWorkspaceSessions);
          } else {
            const newCard: WorkspaceCard = {
              id: event.id,
              name: event.name,
              props: event.props ?? {},
              title: event.name,
              priority: 0,
              sourceCallId: event.id,
            };
            addWsCard(aiMsgId, newCard, setWorkspaceSessions);
          }
        }
      },
      () => {
        // on error
        setIsLoading(false);
      },
      () => {
        // on done
        setIsLoading(false);
        completeWsSession(aiMsgId, setWorkspaceSessions);
      }
    );
  }, [inputValue, isLoading, activeConvId, activeTools, apiModel]);

  const handleRegenerate = useCallback(
    async (msgId: string) => {
      if (!activeConvId || isLoading) return;
      const conv = conversations.find((c) => c.id === activeConvId);
      if (!conv) return;
      const msgIndex = conv.messages.findIndex((m) => m.id === msgId);
      if (msgIndex === -1) return;

      // Find the preceding user message
      const userMsg = conv.messages
        .slice(0, msgIndex)
        .reverse()
        .find((m) => m.role === "user");
      if (!userMsg) return;

      setIsLoading(true);

      // Clear the existing assistant message content for streaming
      const newAiMsgId = `msg-${Date.now()}-regen`;
      setConversations((prev) =>
        prev.map((c) => {
          if (c.id !== activeConvId) return c;
          return {
            ...c,
            messages: c.messages.map((m) =>
              m.id === msgId ? { ...m, content: "", id: newAiMsgId, traceItems: [] } : m
            ),
          };
        })
      );

      const regenConvId = activeConvId;

      abortRef.current = streamChat(
        userMsg.content,
        regenConvId,
        Array.from(activeTools),
        apiModel,
        (event) => {
          if (event.type === "token" && event.content) {
            setConversations((prev) =>
              prev.map((c) => {
                if (c.id !== regenConvId) return c;
                return {
                  ...c,
                  messages: c.messages.map((m) =>
                    m.id === newAiMsgId
                      ? { ...m, content: m.content + event.content }
                      : m
                  ),
                };
              })
            );
          } else if (event.type === "status") {
            // status 事件：确保 assistant 消息占位符已存在（trace 面板早于 token 出现）
            setConversations((prev) =>
              prev.map((c) => {
                if (c.id !== regenConvId) return c;
                const exists = c.messages.some((m) => m.id === newAiMsgId);
                if (exists) return c;
                return {
                  ...c,
                  messages: [
                    ...c.messages,
                    { id: newAiMsgId, role: "assistant" as const, content: "", timestamp: new Date(), traceItems: [] },
                  ],
                };
              })
            );

          } else if (event.type === "tool_start" && event.tool) {
            const item: TraceItem = {
              id: event.call_id ?? `${event.tool}-${event.ts ?? Date.now()}`,
              callId: event.call_id,
              type: "tool_start",
              tool: event.tool,
              title: event.title ?? event.tool,
              inputSummary: event.input_summary,
              itemStatus: "running",
              startTs: event.ts ?? Date.now(),
            };
            appendTraceItem(regenConvId, newAiMsgId, item, setConversations);
          } else if (event.type === "tool_end" && event.tool) {
            updateTraceItem(regenConvId, newAiMsgId, event.call_id, event.tool, {
              type: "tool_end",
              outputSummary: event.output_summary,
              itemStatus: "success",
              durationMs: event.duration_ms,
            }, setConversations);
          } else if (event.type === "tool_error" && event.tool) {
            updateTraceItem(regenConvId, newAiMsgId, event.call_id, event.tool, {
              type: "tool_error",
              errorSummary: event.error_summary,
              itemStatus: "error",
              durationMs: event.duration_ms,
            }, setConversations);
          }
        },
        () => {
          setIsLoading(false);
        },
        () => {
          setIsLoading(false);
        }
      );
    },
    [activeConvId, conversations, isLoading, activeTools, apiModel]
  );

  const handleSuggestionClick = useCallback((text: string) => {
    setInputValue(text);
  }, []);

  return (
    <div className="flex h-screen bg-[var(--gem-bg)] overflow-hidden font-sans" role="main">
      {/* Sidebar */}
      <GeminiSidebar
        isOpen={sidebarOpen}
        conversations={conversations.map(({ id, title }) => ({ id, title }))}
        activeConversationId={activeConvId}
        onNewChat={handleNewChat}
        onSelectConversation={handleSelectConversation}
        onRenameConversation={handleRenameConversation}
        onDeleteConversation={handleDeleteConversation}
        onOpenExtensions={handleOpenExtensions}
        connectedExtensionsCount={extensions.filter((e) => e.connected).length}
      />

      {/* Main content area */}
      <div
        className="flex flex-col flex-1 min-w-0 transition-all duration-300"
        style={{ marginLeft: sidebarOpen ? 260 : 0 }}
      >
        {/* Top bar */}
        <GeminiTopbar
          sidebarOpen={sidebarOpen}
          onToggleSidebar={() => setSidebarOpen((v) => !v)}
        />

        {/* Messages or welcome */}
        <div className="flex-1 overflow-y-auto pt-14 scrollbar-thin">
          {messages.length === 0 && !isLoading ? (
            <>
              <GeminiWelcome onSuggestionClick={handleSuggestionClick} />
              {ticketSystemConnected && ticketSummary && <TicketSummary summary={ticketSummary} />}
            </>
          ) : (
            <div className="max-w-2xl mx-auto px-4 pt-8 pb-4">
              {ticketSystemConnected && ticketSummary && <TicketSummary summary={ticketSummary} />}
              {messages.map((msg) => {
                if (msg.role === "assistant" && msg.content === "" && isLoading) {
                  return null;
                }
                return (
                  <GeminiMessage
                    key={msg.id}
                    message={msg}
                    model={selectedModel}
                    onRegenerate={msg.role === "assistant" ? handleRegenerate : undefined}
                    onFocus={msg.role === "assistant" ? () => setActiveWorkspaceId(msg.id) : undefined}
                  />
                );
              })}
              {isLoading && messages.length > 0 && messages[messages.length - 1].role !== "assistant" && (
                <GeminiTypingIndicator />
              )}
              <div ref={messagesEndRef} />
            </div>
          )}
        </div>

        {/* Input */}
        <GeminiInput
          value={inputValue}
          onChange={setInputValue}
          onSubmit={handleSend}
          isLoading={isLoading}
          activeTools={activeTools}
          onActiveToolsChange={setActiveTools}
          selectedModel={selectedModel}
          onSelectedModelChange={setSelectedModel}
        />
      </div>

      {/* Right: Agent Workspace */}
      <div className="w-[360px] flex-shrink-0 border-l border-[var(--gem-border)] bg-[var(--gem-surface)] overflow-hidden flex flex-col">
        <AgentWorkspace session={activeWorkspace} isGenerating={isLoading} />
      </div>

      {/* Extensions Panel */}
      <ExtensionsPanel
        isOpen={extensionsPanelOpen}
        onClose={() => setExtensionsPanelOpen(false)}
        extensions={extensions}
        onUpdateExtension={handleUpdateExtension}
      />
    </div>
  );
}