blank/frontend-v2/lib/hooks.ts

/**
 * 自定义 React Hooks - 用于与 FastAPI 后端交互
 */
'use client';

import { useState, useCallback, useRef, useEffect } from 'react';
import { apiClient, type ChatMessage, type UserInfo, type UserRole } from './api-client';

// SSE 事件类型
export type SSEEventType =
  | 'start'
  | 'token'
  | 'tools'
  | 'tools_start'
  | 'tool_call'
  | 'tool_done'
  | 'tool_error'
  | 'complete'
  | 'error';

export interface SSEEvent {
  type: SSEEventType;
  data: unknown;
}

// 聊天 Hook
export function useChat() {
  const [isLoading, setIsLoading] = useState(false);
  const [response, setResponse] = useState('');
  const [toolCalls, setToolCalls] = useState<Array<{ tool: string; tool_id?: string; args?: unknown; result: unknown }>>([]);
  const [specs, setSpecs] = useState<any[]>([]);
  const [error, setError] = useState<string | null>(null);
  const abortRef = useRef<(() => void) | null>(null);
  const toolCallsRef = useRef<Array<{ tool: string; result: unknown }>>([]);

  // 保持 ref 同步
  useEffect(() => {
    toolCallsRef.current = toolCalls;
    console.log('[useEffect] toolCalls updated:', toolCalls);
  }, [toolCalls]);

  // 监控 response 状态变化，并在完成时提取 JSON spec
  useEffect(() => {
    console.log('[useEffect] response updated:', {
      length: response.length,
      preview: response.substring(0, 100),
    });

    // 只在加载完成且有响应内容时尝试提取 spec
    // 这确保在流式输出完成后，从 response 中提取 JSON
    if (!isLoading && response && response.trim()) {
      console.log('[useEffect] Loading complete, extracting JSON from response...');
      const extractedData = extractJsonFromMarkdown(response);
      console.log('[useEffect] Extracted data:', extractedData);

      if (extractedData) {
        const newSpecs: any[] = [];

        // 检查是否为有效的组件 spec（有 type 字段）
        if (extractedData.type && typeof extractedData.type === 'string') {
          newSpecs.push(extractedData);
          console.log('[useEffect] Added spec from response:', extractedData);
        }
        // 如果是数组，检查每个元素
        else if (Array.isArray(extractedData)) {
          for (const item of extractedData) {
            if (item && item.type && typeof item.type === 'string') {
              newSpecs.push(item);
              console.log('[useEffect] Added spec from response array:', item);
            }
          }
        }

        if (newSpecs.length > 0) {
          console.log('[useEffect] Setting specs from response:', newSpecs);
          setSpecs(newSpecs);
        }
      }
    }
  }, [response, isLoading]); // extractJsonFromMarkdown 依赖为空数组，不需要在依赖项中

  /**
   * 从 markdown 代码块中提取 JSON
   * 支持格式：```json ... ``` 或 ``` ... ```
   */
  const extractJsonFromMarkdown = useCallback((text: string): any | null => {
    if (typeof text !== 'string') {
      return null;
    }

    // 1. 匹配 ```json ... ``` 代码块（先提取整个代码块内容，不限制格式）
    const jsonCodeBlockMatch = text.match(/```json\s*([\s\S]*?)\s*```/);
    if (jsonCodeBlockMatch) {
      try {
        const jsonStr = jsonCodeBlockMatch[1].trim();
        console.log('[extractJsonFromMarkdown] Found json code block, parsing:', jsonStr.substring(0, 100));
        const parsed = JSON.parse(jsonStr);
        console.log('[extractJsonFromMarkdown] Parsed successfully:', parsed);
        return parsed;
      } catch (e) {
        console.error('[extractJsonFromMarkdown] Failed to parse JSON code block:', e);
        // 继续尝试其他方式
      }
    }

    // 2. 匹配 ``` ... ``` 代码块（无语言标识）
    const codeBlockMatch = text.match(/```\s*([\s\S]*?)\s*```/);
    if (codeBlockMatch) {
      try {
        const jsonStr = codeBlockMatch[1].trim();
        console.log('[extractJsonFromMarkdown] Found code block, trying to parse:', jsonStr.substring(0, 100));
        const parsed = JSON.parse(jsonStr);
        console.log('[extractJsonFromMarkdown] Parsed successfully:', parsed);
        return parsed;
      } catch (e) {
        console.log('[extractJsonFromMarkdown] Code block is not JSON, continuing...');
        // 不是 JSON，继续尝试其他方式
      }
    }

    // 3. 尝试直接解析整个文本
    try {
      console.log('[extractJsonFromMarkdown] Trying to parse entire text as JSON');
      const parsed = JSON.parse(text);
      console.log('[extractJsonFromMarkdown] Parsed successfully:', parsed);
      return parsed;
    } catch {
      console.log('[extractJsonFromMarkdown] Failed to parse as JSON');
      return null;
    }
  }, []);

  // 生成 specs 的通用函数
  const generateSpecs = useCallback(() => {
    const currentToolCalls = toolCallsRef.current;
    console.log('[generateSpecs] Current toolCalls:', currentToolCalls);
    console.log('[generateSpecs] Current response:', response?.substring(0, 200));

    const newSpecs: any[] = [];

    // 1. 从 toolCalls 生成 specs（原有逻辑）
    if (currentToolCalls.length > 0) {
      const { specFromToolCall } = require('@/lib/json-render-catalog');
      const toolSpecs = currentToolCalls.map((call) => specFromToolCall(call.tool, call.result)).filter(Boolean);
      newSpecs.push(...toolSpecs);
      console.log('[generateSpecs] Generated specs from toolCalls:', toolSpecs);
    }

    // 2. 从 response 中提取 JSON 并生成 specs（新增逻辑）
    if (response && response.trim()) {
      const extractedData = extractJsonFromMarkdown(response);
      console.log('[generateSpecs] Extracted JSON from response:', extractedData);

      if (extractedData) {
        // 检查是否为有效的组件 spec（有 type 字段）
        if (extractedData.type && typeof extractedData.type === 'string') {
          newSpecs.push(extractedData);
          console.log('[generateSpecs] Added spec from response:', extractedData);
        }
        // 如果是数组，检查每个元素
        else if (Array.isArray(extractedData)) {
          for (const item of extractedData) {
            if (item && item.type && typeof item.type === 'string') {
              newSpecs.push(item);
              console.log('[generateSpecs] Added spec from response array:', item);
            }
          }
        }
      }
    }

    if (newSpecs.length > 0) {
      console.log('[generateSpecs] Setting specs:', newSpecs);
      setSpecs(newSpecs);
    }

    // 使用 setTimeout 确保 specs 状态更新先被处理
    setTimeout(() => {
      setIsLoading(false);
      console.log('[generateSpecs] Set isLoading to false');
    }, 0);
  }, [response]); // extractJsonFromMarkdown 依赖为空数组，不需要在依赖项中

  const sendMessage = useCallback(async (message: string, history: ChatMessage[] = []) => {
    setIsLoading(true);
    setResponse('');
    setToolCalls([]);
    setSpecs([]);
    setError(null);

    try {
      abortRef.current = await apiClient.chatStreamFetch(
        message,
        history,
        (event) => {
          try {
            console.log('[SSE Raw] type:', event.type, 'data:', event.data?.substring(0, 100));
            const data = JSON.parse(event.data);
            // 将 SSE 事件类型添加到数据中，以便 handleSSEEvent 可以访问
            (data as any).type = event.type;
            handleSSEEvent(data);
          } catch (e) {
            console.error('Failed to parse SSE data:', e);
          }
        },
        (err) => {
          console.error('[sendMessage] Error:', err);
          setError(err.message);
          // 即使出错也生成 specs
          generateSpecs();
        },
        () => {
          // 对话完成时，生成最终的 specs
          generateSpecs();
        }
      );
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Unknown error');
      setIsLoading(false);
    }
  }, [generateSpecs]);

  const handleSSEEvent = (data: unknown) => {
    const event = data as { type?: string; [key: string]: unknown };

    // 调试日志
    console.log('[SSE Event]', event.type, JSON.stringify(event).substring(0, 200));

    switch (event.type) {
      case 'token':
        const tokenData = event as { text?: string };
        if (tokenData.text) {
          setResponse((prev) => {
            const newResponse = prev + tokenData.text;
            console.log('[Token] Adding text:', {
              text: tokenData.text,
              prevLength: prev.length,
              newLength: newResponse.length,
              newPreview: newResponse.substring(0, 100),
            });
            return newResponse;
          });
        }
        break;
      case 'tool_call':
        const toolCallData = event as { tool?: string; args?: unknown; tool_id?: string };
        if (toolCallData.tool) {
          const newToolCall = {
            tool: toolCallData.tool!,
            tool_id: toolCallData.tool_id,
            args: toolCallData.args,  // 保存原始参数
            result: toolCallData.args, // 初始 result 设为 args
          };
          setToolCalls((prev) => [...prev, newToolCall]);
          console.log('[Tool Call] Added:', newToolCall);
        }
        break;
      case 'tool_done':
        const toolDoneData = event as { tool?: string; tool_id?: string; result?: unknown };
        console.log('[Tool Done] Raw event:', toolDoneData);
        console.log('[Tool Done] Result type:', typeof toolDoneData.result, 'Result:', toolDoneData.result);
        // 优先使用 tool_id 匹配，回退到 tool 名称匹配
        if (toolDoneData.tool_id || toolDoneData.tool) {
          setToolCalls((prev) => {
            console.log('[Tool Done] Current toolCalls before update:', prev);
            console.log('[Tool Done] Looking for tool_id:', toolDoneData.tool_id, 'type:', typeof toolDoneData.tool_id);
            const updated = [...prev];
            // 使用 tool_id 查找
            let index = updated.findIndex((t) => {
              console.log('[Tool Done] Comparing:', t.tool_id, '===', toolDoneData.tool_id, 'result:', t.tool_id === toolDoneData.tool_id);
              return t.tool_id === toolDoneData.tool_id;
            });
            console.log('[Tool Done] Found index:', index);
            // 回退到 tool 名称查找
            if (index < 0 && toolDoneData.tool) {
              index = updated.findIndex((t) => t.tool === toolDoneData.tool);
              console.log('[Tool Done] Found by tool name index:', index);
            }
            if (index >= 0) {
              // 保留 args，只更新 result
              const existing = updated[index];
              updated[index] = {
                ...existing,
                result: toolDoneData.result,
                // 确保 args 不被覆盖
                args: existing.args,
              };
              console.log('[Tool Done] Updated tool call:', updated[index]);
            } else {
              console.log('[Tool Done] No matching tool_call found for:', toolDoneData);
            }
            return updated;
          });
        }
        break;
      case 'complete':
        const completeData = event as { response?: string; tool_calls?: unknown };
        console.log('[Complete]', completeData);
        console.log('[Complete] Response field:', completeData.response?.substring(0, 100));
        // 使用函数式更新来避免闭包陷阱
        setResponse((prev) => {
          // 如果 complete 事件有响应，优先使用它（因为它包含完整响应）
          // 否则保留从 token 事件构建的响应
          if (completeData.response) {
            console.log('[Complete] Using complete response, replacing prev with length:', prev.length);
            return completeData.response;
          }
          console.log('[Complete] No response in complete event, keeping prev with length:', prev.length);
          return prev;
        });
        // 不在这里设置 isLoading(false)，让 onComplete 回调处理
        // 这样可以确保 specs 在 isLoading 变为 false 之前生成
        break;
      case 'error':
        const errorData = event as { error?: string };
        console.error('[SSE Error]', errorData);
        setError(errorData.error || 'Unknown error');
        // 不在这里设置 isLoading(false)，让 stream 的 onError/onComplete 处理
        // 这样可以确保 specs 在 isLoading 变为 false 之前生成
        break;
    }
  };

  const abort = useCallback(() => {
    abortRef.current?.();
    setIsLoading(false);
  }, []);

  return {
    sendMessage,
    abort,
    isLoading,
    response,
    toolCalls,
    specs,
    error,
  };
}

// 认证 Hook
export function useAuth() {
  const [isAuthenticated, setIsAuthenticated] = useState(false);
  const [username, setUsername] = useState<string | null>(null);
  const [role, setRole] = useState<UserRole | null>(null);
  const [isLoading, setIsLoading] = useState(false);

  // 从 localStorage 恢复认证状态
  useEffect(() => {
    const storedSessionId = localStorage.getItem('session_id');
    const storedUserInfo = localStorage.getItem('userInfo');
    if (storedSessionId) {
      apiClient.setSession(storedSessionId);
      setIsAuthenticated(true);
    }
    if (storedUserInfo) {
      try {
        const userInfo: UserInfo = JSON.parse(storedUserInfo);
        setUsername(userInfo.username);
        setRole(userInfo.role);
        apiClient.setUserInfo(userInfo);
      } catch (e) {
        console.error('Failed to parse stored userInfo:', e);
      }
    }
  }, []);

  // 根据角色获取 MCP URL
  const getMcpUrl = useCallback((): string => {
    return apiClient.getMcpUrl();
  }, []);

  const login = useCallback(async (email: string, password: string) => {
    setIsLoading(true);
    try {
      console.log('[useAuth] Starting login:', email);
      const response = await apiClient.login(email, password);
      console.log('[useAuth] Login response:', response);
      // 检查 session_id 而非 success 字段
      if (response.session_id) {
        console.log('[useAuth] Login successful, setting state...');
        setIsAuthenticated(true);
        setUsername(response.username);
        setRole(response.role);
        // 持久化到 localStorage
        localStorage.setItem('session_id', response.session_id);
        localStorage.setItem('username', response.username);
        const userInfo: UserInfo = {
          username: response.username,
          role: response.role,
        };
        localStorage.setItem('userInfo', JSON.stringify(userInfo));
        console.log('[useAuth] State updated, userInfo:', userInfo);
        return response;
      }
      throw new Error('Login failed: No session_id returned');
    } finally {
      setIsLoading(false);
    }
  }, []);

  const register = useCallback(async (email: string, username: string, password: string) => {
    setIsLoading(true);
    try {
      // 先注册
      const registerResponse = await apiClient.register(email, username, password);
      if (!registerResponse.success) {
        throw new Error(registerResponse.message || 'Registration failed');
      }
      // 注册成功后自动登录
      const loginResponse = await apiClient.login(email, password);
      if (loginResponse.session_id) {
        setIsAuthenticated(true);
        setUsername(loginResponse.username);
        setRole(loginResponse.role);
        // 持久化到 localStorage
        localStorage.setItem('session_id', loginResponse.session_id);
        localStorage.setItem('username', loginResponse.username);
        const userInfo: UserInfo = {
          username: loginResponse.username,
          role: loginResponse.role,
        };
        localStorage.setItem('userInfo', JSON.stringify(userInfo));
        return { ...registerResponse, session: loginResponse };
      }
      throw new Error('Auto-login after registration failed');
    } finally {
      setIsLoading(false);
    }
  }, []);

  const logout = useCallback(async () => {
    setIsLoading(true);
    try {
      await apiClient.logout();
      setIsAuthenticated(false);
      setUsername(null);
      setRole(null);
      // 清除 localStorage
      localStorage.removeItem('session_id');
      localStorage.removeItem('username');
      localStorage.removeItem('userInfo');
    } finally {
      setIsLoading(false);
    }
  }, []);

  const checkStatus = useCallback(async () => {
    try {
      const status = await apiClient.authStatus();
      setIsAuthenticated(status.authenticated);
      setUsername(status.username || null);
      setRole((status.role as UserRole) || null);
      return status;
    } catch (err) {
      setIsAuthenticated(false);
      setUsername(null);
      setRole(null);
      return null;
    }
  }, []);

  return {
    isAuthenticated,
    username,
    role,
    isLoading,
    login,
    register,
    logout,
    checkStatus,
    getMcpUrl,
  };
}

// MCP 工具 Hook
export function useMcpTools() {
  const [tools, setTools] = useState<Array<{ name: string; description: string }>>([]);
  const [isLoading, setIsLoading] = useState(false);

  const fetchTools = useCallback(async () => {
    setIsLoading(true);
    try {
      const response = await apiClient.listMcpTools();
      setTools(response.tools);
      return response;
    } finally {
      setIsLoading(false);
    }
  }, []);

  return {
    tools,
    isLoading,
    fetchTools,
  };
}