zclaw_openfang/desktop/src/store/chatStore.ts

import { create } from 'zustand';
import { persist } from 'zustand/middleware';
import { getGatewayClient, AgentStreamDelta } from '../lib/gateway-client';
import { intelligenceClient } from '../lib/intelligence-client';
import { getMemoryExtractor } from '../lib/memory-extractor';
import { getAgentSwarm } from '../lib/agent-swarm';
import { getSkillDiscovery } from '../lib/skill-discovery';
import { useOfflineStore, isOffline } from './offlineStore';
import { useConnectionStore } from './connectionStore';

export interface MessageFile {
  name: string;
  path?: string;
  size?: number;
  type?: string;
}

export interface CodeBlock {
  language?: string;
  filename?: string;
  content?: string;
}

export interface Message {
  id: string;
  role: 'user' | 'assistant' | 'tool' | 'hand' | 'workflow' | 'system';
  content: string;
  timestamp: Date;
  runId?: string;
  streaming?: boolean;
  toolName?: string;
  toolInput?: string;
  toolOutput?: string;
  error?: string;
  // Hand event fields
  handName?: string;
  handStatus?: string;
  handResult?: unknown;
  // Workflow event fields
  workflowId?: string;
  workflowStep?: string;
  workflowStatus?: string;
  workflowResult?: unknown;
  // Output files and code blocks
  files?: MessageFile[];
  codeBlocks?: CodeBlock[];
}

export interface Conversation {
  id: string;
  title: string;
  messages: Message[];
  sessionKey: string | null;
  agentId: string | null;
  createdAt: Date;
  updatedAt: Date;
}

export interface Agent {
  id: string;
  name: string;
  icon: string;
  color: string;
  lastMessage: string;
  time: string;
}

export interface AgentProfileLike {
  id: string;
  name: string;
  nickname?: string;
  role?: string;
}

interface ChatState {
  messages: Message[];
  conversations: Conversation[];
  currentConversationId: string | null;
  agents: Agent[];
  currentAgent: Agent | null;
  isStreaming: boolean;
  isLoading: boolean;
  currentModel: string;
  sessionKey: string | null;

  addMessage: (message: Message) => void;
  updateMessage: (id: string, updates: Partial<Message>) => void;
  setIsLoading: (loading: boolean) => void;
  setCurrentAgent: (agent: Agent) => void;
  syncAgents: (profiles: AgentProfileLike[]) => void;
  setCurrentModel: (model: string) => void;
  sendMessage: (content: string) => Promise<void>;
  initStreamListener: () => () => void;
  newConversation: () => void;
  switchConversation: (id: string) => void;
  deleteConversation: (id: string) => void;
  dispatchSwarmTask: (description: string, style?: 'sequential' | 'parallel' | 'debate') => Promise<string | null>;
  searchSkills: (query: string) => { results: Array<{ id: string; name: string; description: string }>; totalAvailable: number };
}

function generateConvId(): string {
  return `conv_${Date.now()}_${Math.random().toString(36).slice(2, 6)}`;
}

function deriveTitle(messages: Message[]): string {
  const firstUser = messages.find(m => m.role === 'user');
  if (firstUser) {
    const text = firstUser.content.trim();
    return text.length > 30 ? text.slice(0, 30) + '...' : text;
  }
  return '新对话';
}

const DEFAULT_AGENT: Agent = {
  id: '1',
  name: 'ZCLAW',
  icon: '🦞',
  color: 'bg-gradient-to-br from-orange-500 to-red-500',
  lastMessage: '发送消息开始对话',
  time: '',
};

export function toChatAgent(profile: AgentProfileLike): Agent {
  return {
    id: profile.id,
    name: profile.name,
    icon: profile.nickname?.slice(0, 1) || '🦞',
    color: 'bg-gradient-to-br from-orange-500 to-red-500',
    lastMessage: profile.role || '新分身',
    time: '',
  };
}

function resolveConversationAgentId(agent: Agent | null): string | null {
  if (!agent || agent.id === DEFAULT_AGENT.id) {
    return null;
  }
  return agent.id;
}

function resolveGatewayAgentId(agent: Agent | null): string | undefined {
  if (!agent || agent.id === DEFAULT_AGENT.id || agent.id.startsWith('clone_')) {
    return undefined;
  }
  return agent.id;
}

function resolveAgentForConversation(agentId: string | null, agents: Agent[]): Agent {
  if (!agentId) {
    return DEFAULT_AGENT;
  }
  return agents.find((agent) => agent.id === agentId) || DEFAULT_AGENT;
}

function upsertActiveConversation(
  conversations: Conversation[],
  state: Pick<ChatState, 'messages' | 'sessionKey' | 'currentConversationId' | 'currentAgent'>
): Conversation[] {
  if (state.messages.length === 0) {
    return conversations;
  }

  const currentId = state.currentConversationId || generateConvId();
  const existingIdx = conversations.findIndex((conversation) => conversation.id === currentId);
  const nextConversation: Conversation = {
    id: currentId,
    title: deriveTitle(state.messages),
    messages: [...state.messages],
    sessionKey: state.sessionKey,
    agentId: resolveConversationAgentId(state.currentAgent),
    createdAt: existingIdx >= 0 ? conversations[existingIdx].createdAt : new Date(),
    updatedAt: new Date(),
  };

  if (existingIdx >= 0) {
    conversations[existingIdx] = nextConversation;
    return conversations;
  }

  return [nextConversation, ...conversations];
}

export const useChatStore = create<ChatState>()(
  persist(
    (set, get) => ({
  messages: [],
  conversations: [],
  currentConversationId: null,
  agents: [DEFAULT_AGENT],
  currentAgent: DEFAULT_AGENT,
  isStreaming: false,
  isLoading: false,
  currentModel: 'glm-5',
  sessionKey: null,

  addMessage: (message) =>
    set((state) => ({ messages: [...state.messages, message] })),

  updateMessage: (id, updates) =>
    set((state) => ({
      messages: state.messages.map((m) =>
        m.id === id ? { ...m, ...updates } : m
      ),
    })),

  setIsLoading: (loading) => set({ isLoading: loading }),

  setCurrentAgent: (agent) =>
    set((state) => {
      if (state.currentAgent?.id === agent.id) {
        return { currentAgent: agent };
      }

      // Save current conversation before switching
      const conversations = upsertActiveConversation([...state.conversations], state);

      // Try to find existing conversation for this agent
      const agentConversation = conversations.find(c => c.agentId === agent.id);

      if (agentConversation) {
        // Restore the agent's previous conversation
        return {
          conversations,
          currentAgent: agent,
          messages: [...agentConversation.messages],
          sessionKey: agentConversation.sessionKey,
          isStreaming: false,
          currentConversationId: agentConversation.id,
        };
      }

      // No existing conversation, start fresh
      return {
        conversations,
        currentAgent: agent,
        messages: [],
        sessionKey: null,
        isStreaming: false,
        currentConversationId: null,
      };
    }),

  syncAgents: (profiles) =>
    set((state) => {
      const agents = profiles.length > 0 ? profiles.map(toChatAgent) : [DEFAULT_AGENT];
      const currentAgent = state.currentConversationId
        ? resolveAgentForConversation(
            state.conversations.find((conversation) => conversation.id === state.currentConversationId)?.agentId || null,
            agents
          )
        : state.currentAgent
          ? agents.find((agent) => agent.id === state.currentAgent?.id) || agents[0]
          : agents[0];
      return { agents, currentAgent };
    }),

  setCurrentModel: (model) => set({ currentModel: model }),

  newConversation: () => {
    const state = get();
    const conversations = upsertActiveConversation([...state.conversations], state);

    set({
      conversations,
      messages: [],
      sessionKey: null,
      isStreaming: false,
      currentConversationId: null,
    });
  },

  switchConversation: (id: string) => {
    const state = get();
    const conversations = upsertActiveConversation([...state.conversations], state);

    const target = conversations.find(c => c.id === id);
    if (target) {
      set({
        conversations,
        messages: [...target.messages],
        sessionKey: target.sessionKey,
        currentAgent: resolveAgentForConversation(target.agentId, state.agents),
        currentConversationId: target.id,
        isStreaming: false,
      });
    }
  },

  deleteConversation: (id: string) => {
    const state = get();
    const conversations = state.conversations.filter(c => c.id !== id);
    if (state.currentConversationId === id) {
      set({ conversations, messages: [], sessionKey: null, currentConversationId: null, isStreaming: false });
    } else {
      set({ conversations });
    }
  },

  sendMessage: async (content: string) => {
    const { addMessage, currentAgent, sessionKey } = get();
    const effectiveSessionKey = sessionKey || `session_${Date.now()}`;
    const effectiveAgentId = resolveGatewayAgentId(currentAgent);
    const agentId = currentAgent?.id || 'zclaw-main';

    // Check if offline - queue message instead of sending
    if (isOffline()) {
      const { queueMessage } = useOfflineStore.getState();
      const queueId = queueMessage(content, effectiveAgentId, effectiveSessionKey);
      console.log(`[Chat] Offline - message queued: ${queueId}`);

      // Show a system message about offline queueing
      const systemMsg: Message = {
        id: `system_${Date.now()}`,
        role: 'system',
        content: `后端服务不可用，消息已保存到本地队列。恢复连接后将自动发送。`,
        timestamp: new Date(),
      };
      addMessage(systemMsg);

      // Add user message for display
      const userMsg: Message = {
        id: `user_${Date.now()}`,
        role: 'user',
        content,
        timestamp: new Date(),
      };
      addMessage(userMsg);
      return;
    }

    // Check context compaction threshold before adding new message
    try {
      const messages = get().messages.map(m => ({ role: m.role, content: m.content }));
      const check = await intelligenceClient.compactor.checkThreshold(messages);
      if (check.should_compact) {
        console.log(`[Chat] Context compaction triggered (${check.urgency}): ${check.current_tokens} tokens`);
        const result = await intelligenceClient.compactor.compact(
          get().messages.map(m => ({
            role: m.role,
            content: m.content,
            id: m.id,
            timestamp: m.timestamp instanceof Date ? m.timestamp.toISOString() : m.timestamp
          })),
          agentId,
          get().currentConversationId ?? undefined
        );
        // Replace messages with compacted version
        const compactedMsgs: Message[] = result.compacted_messages.map((m, i) => ({
          id: m.id || `compacted_${i}_${Date.now()}`,
          role: m.role as Message['role'],
          content: m.content,
          timestamp: m.timestamp ? new Date(m.timestamp) : new Date(),
        }));
        set({ messages: compactedMsgs });
      }
    } catch (err) {
      console.warn('[Chat] Context compaction check failed:', err);
    }

    // Build memory-enhanced content
    let enhancedContent = content;
    try {
      const relevantMemories = await intelligenceClient.memory.search({
        agentId,
        query: content,
        limit: 8,
        minImportance: 3,
      });
      const memoryContext = relevantMemories.length > 0
        ? `\n\n## 相关记忆\n${relevantMemories.map(m => `- [${m.type}] ${m.content}`).join('\n')}`
        : '';
      const systemPrompt = await intelligenceClient.identity.buildPrompt(agentId, memoryContext);
      if (systemPrompt) {
        enhancedContent = `<context>\n${systemPrompt}\n</context>\n\n${content}`;
      }
    } catch (err) {
      console.warn('[Chat] Memory enhancement failed, proceeding without:', err);
    }

    // Add user message (original content for display)
    const userMsg: Message = {
      id: `user_${Date.now()}`,
      role: 'user',
      content,
      timestamp: new Date(),
    };
    addMessage(userMsg);

    // Create placeholder assistant message for streaming
    const assistantId = `assistant_${Date.now()}`;
    const assistantMsg: Message = {
      id: assistantId,
      role: 'assistant',
      content: '',
      timestamp: new Date(),
      streaming: true,
    };
    addMessage(assistantMsg);
    set({ isStreaming: true });

    try {
      const client = getGatewayClient();

      // Check connection state first
      const connectionState = useConnectionStore.getState().connectionState;

      if (connectionState !== 'connected') {
        // Connection lost during send - update error
        throw new Error(`Not connected (state: ${connectionState})`);
      }

      // Try streaming first (OpenFang WebSocket)
      const { runId } = await client.chatStream(
        enhancedContent,
        {
          onDelta: () => { /* Handled by initStreamListener to prevent duplication */ },
          onTool: (tool: string, input: string, output: string) => {
            const toolMsg: Message = {
              id: `tool_${Date.now()}_${Math.random().toString(36).slice(2, 6)}`,
              role: 'tool',
              content: output || input,
              timestamp: new Date(),
              runId,
              toolName: tool,
              toolInput: input,
              toolOutput: output,
            };
            set((state) => ({ messages: [...state.messages, toolMsg] }));
          },
          onHand: (name: string, status: string, result?: unknown) => {
            const handMsg: Message = {
              id: `hand_${Date.now()}_${Math.random().toString(36).slice(2, 6)}`,
              role: 'hand',
              content: result
                ? (typeof result === 'string' ? result : JSON.stringify(result, null, 2))
                : `Hand: ${name} - ${status}`,
              timestamp: new Date(),
              runId,
              handName: name,
              handStatus: status,
              handResult: result,
            };
            set((state) => ({ messages: [...state.messages, handMsg] }));
          },
          onComplete: () => {
            const state = get();

            // Save conversation to persist across refresh
            const conversations = upsertActiveConversation([...state.conversations], state);
            const currentConvId = state.currentConversationId || conversations[0]?.id;

            set({
              isStreaming: false,
              conversations,
              currentConversationId: currentConvId,
              messages: state.messages.map((m) =>
                m.id === assistantId ? { ...m, streaming: false, runId } : m
              ),
            });

            // Async memory extraction after stream completes
            const msgs = get().messages
              .filter(m => m.role === 'user' || m.role === 'assistant')
              .map(m => ({ role: m.role, content: m.content }));
            getMemoryExtractor().extractFromConversation(msgs, agentId, get().currentConversationId ?? undefined).catch(err => {
              console.warn('[Chat] Memory extraction failed:', err);
            });
            // Track conversation for reflection trigger
            intelligenceClient.reflection.recordConversation().catch(err => {
              console.warn('[Chat] Recording conversation failed:', err);
            });
            intelligenceClient.reflection.shouldReflect().then(shouldReflect => {
              if (shouldReflect) {
                intelligenceClient.reflection.reflect(agentId, []).catch(err => {
                  console.warn('[Chat] Reflection failed:', err);
                });
              }
            });
          },
          onError: (error: string) => {
            set((state) => ({
              isStreaming: false,
              messages: state.messages.map((m) =>
                m.id === assistantId
                  ? { ...m, content: `⚠️ ${error}`, streaming: false, error }
                  : m
              ),
            }));
          },
        },
        {
          sessionKey: effectiveSessionKey,
          agentId: effectiveAgentId,
        }
      );

      if (!sessionKey) {
        set({ sessionKey: effectiveSessionKey });
      }

      // Store runId on the message for correlation
      set((state) => ({
        messages: state.messages.map((m) =>
          m.id === assistantId ? { ...m, runId } : m
        ),
      }));
    } catch (err: unknown) {
      // Gateway not connected — show error in the assistant bubble
      const errorMessage = err instanceof Error ? err.message : '无法连接 Gateway';
      set((state) => ({
        isStreaming: false,
        messages: state.messages.map((m) =>
          m.id === assistantId
            ? {
                ...m,
                content: `⚠️ ${errorMessage}`,
                streaming: false,
                error: errorMessage,
              }
            : m
        ),
      }));
    }
  },

  dispatchSwarmTask: async (description: string, style?: 'sequential' | 'parallel' | 'debate') => {
    try {
      const swarm = getAgentSwarm();
      const task = swarm.createTask(description, {
        communicationStyle: style || 'parallel',
      });

      // Set up executor that uses gateway client
      swarm.setExecutor(async (agentId: string, prompt: string, context?: string) => {
        const client = getGatewayClient();
        const fullPrompt = context ? `${context}\n\n${prompt}` : prompt;
        const result = await client.chat(fullPrompt, { agentId: agentId.startsWith('clone_') ? undefined : agentId });
        return result?.response || '(无响应)';
      });

      const result = await swarm.execute(task);

      // Add swarm result as assistant message
      const swarmMsg: Message = {
        id: `swarm_${Date.now()}`,
        role: 'assistant',
        content: result.summary || '协作任务完成',
        timestamp: new Date(),
      };
      get().addMessage(swarmMsg);

      return result.task.id;
    } catch (err) {
      console.warn('[Chat] Swarm dispatch failed:', err);
      return null;
    }
  },

  searchSkills: (query: string) => {
    const discovery = getSkillDiscovery();
    const result = discovery.searchSkills(query);
    return {
      results: result.results.map(s => ({ id: s.id, name: s.name, description: s.description })),
      totalAvailable: result.totalAvailable,
    };
  },

  initStreamListener: () => {
    const client = getGatewayClient();

    const unsubscribe = client.onAgentStream((delta: AgentStreamDelta) => {
      const state = get();

      const streamingMsg = [...state.messages]
        .reverse()
        .find((m) => (
          m.role === 'assistant'
          && m.streaming
          && (
            (delta.runId && m.runId === delta.runId)
            || (!delta.runId && m.runId == null)
          )
        ))
        || [...state.messages]
          .reverse()
          .find((m) => m.role === 'assistant' && m.streaming);

      if (!streamingMsg) return;

      if (delta.stream === 'assistant' && (delta.delta || delta.content)) {
        set((s) => ({
          messages: s.messages.map((m) =>
            m.id === streamingMsg.id
              ? { ...m, content: m.content + (delta.delta || delta.content || '') }
              : m
          ),
        }));
      } else if (delta.stream === 'tool') {
        const toolMsg: Message = {
          id: `tool_${Date.now()}_${Math.random().toString(36).slice(2, 6)}`,
          role: 'tool',
          content: delta.toolOutput || '',
          timestamp: new Date(),
          runId: delta.runId,
          toolName: delta.tool,
          toolInput: delta.toolInput,
          toolOutput: delta.toolOutput,
        };
        set((s) => ({ messages: [...s.messages, toolMsg] }));
      } else if (delta.stream === 'lifecycle') {
        if (delta.phase === 'end' || delta.phase === 'error') {
          set((s) => ({
            isStreaming: false,
            messages: s.messages.map((m) =>
              m.id === streamingMsg.id
                ? {
                    ...m,
                    streaming: false,
                    error: delta.phase === 'error' ? delta.error : undefined,
                  }
                : m
            ),
          }));
        }
      } else if (delta.stream === 'hand') {
        // Handle Hand trigger events from OpenFang
        const handMsg: Message = {
          id: `hand_${Date.now()}_${Math.random().toString(36).slice(2, 6)}`,
          role: 'hand',
          content: delta.handResult
            ? (typeof delta.handResult === 'string' ? delta.handResult : JSON.stringify(delta.handResult, null, 2))
            : `Hand: ${delta.handName || 'unknown'} - ${delta.handStatus || 'triggered'}`,
          timestamp: new Date(),
          runId: delta.runId,
          handName: delta.handName,
          handStatus: delta.handStatus,
          handResult: delta.handResult,
        };
        set((s) => ({ messages: [...s.messages, handMsg] }));
      } else if (delta.stream === 'workflow') {
        // Handle Workflow execution events from OpenFang
        const workflowMsg: Message = {
          id: `workflow_${Date.now()}_${Math.random().toString(36).slice(2, 6)}`,
          role: 'workflow',
          content: delta.workflowResult
            ? (typeof delta.workflowResult === 'string' ? delta.workflowResult : JSON.stringify(delta.workflowResult, null, 2))
            : `Workflow: ${delta.workflowId || 'unknown'} step ${delta.workflowStep || '?'} - ${delta.workflowStatus || 'running'}`,
          timestamp: new Date(),
          runId: delta.runId,
          workflowId: delta.workflowId,
          workflowStep: delta.workflowStep,
          workflowStatus: delta.workflowStatus,
          workflowResult: delta.workflowResult,
        };
        set((s) => ({ messages: [...s.messages, workflowMsg] }));
      }
    });

    return unsubscribe;
  },
}),
    {
      name: 'zclaw-chat-storage',
      partialize: (state) => ({
        conversations: state.conversations,
        currentModel: state.currentModel,
        currentAgentId: state.currentAgent?.id,
        currentConversationId: state.currentConversationId,
      }),
      onRehydrateStorage: () => (state) => {
        // Rehydrate Date objects from JSON strings
        if (state?.conversations) {
          for (const conv of state.conversations) {
            conv.createdAt = new Date(conv.createdAt);
            conv.updatedAt = new Date(conv.updatedAt);
            for (const msg of conv.messages) {
              msg.timestamp = new Date(msg.timestamp);
              msg.streaming = false; // Never restore streaming state
            }
          }
        }

        // Restore messages from current conversation if exists
        if (state?.currentConversationId && state.conversations) {
          const currentConv = state.conversations.find(c => c.id === state.currentConversationId);
          if (currentConv) {
            state.messages = [...currentConv.messages];
            state.sessionKey = currentConv.sessionKey;
          }
        }
      },
    },
  ),
);

// Dev-only: Expose chatStore to window for E2E testing
if (import.meta.env.DEV && typeof window !== 'undefined') {
  (window as any).__ZCLAW_STORES__ = (window as any).__ZCLAW_STORES__ || {};
  (window as any).__ZCLAW_STORES__.chat = useChatStore;
}