platform/src/workspaces/apps/Chat/usePi.ts

import { useState, useEffect, useRef } from 'react';
import { useChatWebSocket } from 'hooks/useChatWebSocket';
import { useChatSessions } from '@/state/useChatSessions';
import type { ChatMessage, ServerMessage, TaskInfo, Message } from './types';

const SAVE_DEBOUNCE_MS = 1000;

type ResourceChatStorage = {
  load: () => Promise<{ sessionId: string | null; messages: ChatMessage[] }>;
  save: (sessionId: string, messages: ChatMessage[]) => Promise<void>;
};

type UsePiOptions = {
  replaceUrl?: boolean;
  storage?: ResourceChatStorage;
  resourceChatDir?: string;
  taskInfo?: TaskInfo;
};

export function usePi(initialSessionId?: string, initialModel?: string | null, options?: UsePiOptions) {
  const { replaceUrl = true, storage, resourceChatDir, taskInfo } = options ?? {};
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [streamingText, setStreamingText] = useState('');
  const [isGenerating, setIsGenerating] = useState(false);
  const [sessionId, setSessionId] = useState<string | null>(initialSessionId ?? null);
  const [model, setModel] = useState<string | null>(null);
  const [selectedModel, setSelectedModel] = useState<string | null>(initialModel ?? null);
  const [cwd, setCwd] = useState<string | null>(null);

  const streamingRef = useRef('');
  const rafRef = useRef<number | null>(null);
  const sessionIdRef = useRef<string | null>(initialSessionId ?? null);
  const saveTimerRef = useRef<number | null>(null);

  const { getSession, saveMessages } = useChatSessions();

  const token = localStorage.getItem('BEARER_TOKEN');
  const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
  const wsUrl = `${protocol}//${window.location.host}/api/pi/chat/ws?token=${token}`;

  function flushStreaming() {
    if (rafRef.current !== null) cancelAnimationFrame(rafRef.current);
    rafRef.current = requestAnimationFrame(() => {
      setStreamingText(streamingRef.current);
      rafRef.current = null;
    });
  }

  function commitStreaming() {
    if (!streamingRef.current) return;
    setMessages((prev) => [...prev, { role: 'assistant', text: streamingRef.current }]);
    streamingRef.current = '';
    setStreamingText('');
  }

  function handleMessage(data: unknown) {
    const msg = data as ServerMessage;

    switch (msg.type) {
      case 'session:init':
        sessionIdRef.current = msg.sessionId;
        setSessionId(msg.sessionId);
        setModel(msg.model);
        setCwd(msg.cwd);
        if (replaceUrl) window.history.replaceState(null, '', `/chat/${msg.sessionId}`);
        break;

      case 'assistant:delta':
        streamingRef.current += msg.text;
        flushStreaming();
        break;

      case 'assistant:text':
        if (streamingRef.current) {
          commitStreaming();
        } else {
          setMessages((prev) => [...prev, { role: 'assistant', text: msg.text }]);
        }
        break;

      case 'tool:start':
        setMessages((prev) => [
          ...prev,
          {
            role: 'tool',
            toolName: msg.toolName,
            toolInput: msg.toolInput,
            toolCallId: msg.toolCallId,
          },
        ]);
        break;

      case 'tool:result':
        setMessages((prev) =>
          prev.map((m) =>
            m.role === 'tool' && m.toolCallId === msg.toolCallId
              ? { ...m, output: msg.output, isError: msg.isError }
              : m,
          ),
        );
        break;

      case 'result':
        commitStreaming();
        setMessages((prev) => [
          ...prev,
          {
            role: 'result',
            cost: msg.cost,
          },
        ]);
        setIsGenerating(false);
        break;

      case 'sync:messages':
        sessionIdRef.current = msg.sessionId;
        setSessionId(msg.sessionId);
        // Convert Message[] to ChatMessage[]
        const chatMessages = msg.messages.map((m): ChatMessage => {
          if (m.role === 'user') {
            return { role: 'user', text: m.text || '' };
          } else if (m.role === 'assistant') {
            return { role: 'assistant', text: m.text || '' };
          } else if (m.role === 'tool') {
            return {
              role: 'tool',
              toolName: m.toolName || '',
              toolInput: m.toolInput || {},
              toolCallId: m.toolCallId || '',
              output: m.output,
              isError: m.isError,
            };
          }
          return { role: 'assistant', text: '' }; // Fallback
        });
        setMessages(chatMessages);
        setIsGenerating(msg.isGenerating);
        if (msg.streamingText) {
          streamingRef.current = msg.streamingText;
          flushStreaming();
        }
        break;

      case 'error':
        commitStreaming();
        setMessages((prev) => [...prev, { role: 'error', text: msg.message }]);
        setIsGenerating(false);
        break;

      case 'stopped':
        commitStreaming();
        setIsGenerating(false);
        break;
    }
  }

  const { isConnected, send } = useChatWebSocket({ url: wsUrl, onMessage: handleMessage });

  // Load messages from server on mount when resuming a session
  useEffect(() => {
    if (storage) {
      storage
        .load()
        .then(({ sessionId: sid, messages: msgs }) => {
          if (sid) {
            sessionIdRef.current = sid;
            setSessionId(sid);
          }
          if (msgs.length > 0) setMessages(msgs);
        })
        .catch(() => {});
      return;
    }
    if (!initialSessionId) return;
    getSession(initialSessionId)
      .then((data) => {
        if (data.session?.messages && data.session.messages.length > 0) {
          // Convert backend Message[] to ChatMessage[]
          const chatMessages = data.session.messages.map((m: Message): ChatMessage => {
            if (m.role === 'user') {
              return { role: 'user', text: m.text || '' };
            } else if (m.role === 'assistant') {
              return { role: 'assistant', text: m.text || '' };
            } else if (m.role === 'tool') {
              return {
                role: 'tool',
                toolName: m.toolName || '',
                toolInput: m.toolInput || {},
                toolCallId: m.toolCallId || '',
                output: m.output,
                isError: m.isError,
              };
            }
            return { role: 'assistant', text: '' }; // Fallback
          });
          setMessages(chatMessages);
        }
      })
      .catch(() => {});
  }, [initialSessionId]);

  // Debounced save messages to server
  useEffect(() => {
    if (!sessionIdRef.current || messages.length === 0) return;

    if (saveTimerRef.current !== null) clearTimeout(saveTimerRef.current);

    const sid = sessionIdRef.current;
    const snapshot = messages;
    saveTimerRef.current = window.setTimeout(() => {
      if (storage) {
        storage.save(sid, snapshot).catch(() => {});
      } else {
        saveMessages(sid, snapshot).catch(() => {});
      }
      saveTimerRef.current = null;
    }, SAVE_DEBOUNCE_MS);

    return () => {
      if (saveTimerRef.current !== null) {
        clearTimeout(saveTimerRef.current);
        saveTimerRef.current = null;
      }
    };
  }, [messages]);

  // Clean up RAF on unmount
  useEffect(() => {
    return () => {
      if (rafRef.current !== null) cancelAnimationFrame(rafRef.current);
    };
  }, []);

  function sendPrompt(
    text: string,
    attachmentIds?: string[],
    images?: { filename: string; dataUrl: string }[],
    cwdParam?: { root?: string; path: string },
    groupSlug?: string | null,
  ) {
    setMessages((prev) => [...prev, { role: 'user', text, ...(images?.length ? { images } : {}) }]);
    setIsGenerating(true);
    streamingRef.current = '';
    setStreamingText('');

    // Parse dataUrls into { mediaType, data } for the server
    const imageData = images
      ?.map((img) => {
        const match = img.dataUrl.match(/^data:([^;]+);base64,(.+)$/);
        return match ? { mediaType: match[1], data: match[2] } : null;
      })
      .filter((x): x is { mediaType: string; data: string } => x !== null);

    send({
      type: 'chat',
      prompt: text,
      sessionId: sessionIdRef.current,
      ...(selectedModel ? { model: selectedModel } : {}),
      ...(cwdParam?.path ? { cwd: cwdParam.path } : {}),
      ...(groupSlug !== undefined ? { groupSlug } : {}),
      ...(attachmentIds?.length ? { attachmentIds } : {}),
      ...(imageData?.length ? { images: imageData } : {}),
      ...(resourceChatDir ? { resourceChatDir } : {}),
      ...(taskInfo ? { taskInfo } : {}),
    });
  }

  function stopGeneration() {
    send({ type: 'stop' });
  }

  return {
    messages,
    streamingText,
    isConnected,
    isGenerating,
    sessionId,
    model,
    selectedModel,
    cwd,
    setSelectedModel,
    sendPrompt,
    stopGeneration,
  };
}

export type UsePiType = ReturnType<typeof usePi>;