platform/src/workspaces/plugins/Chat/client/useOpenCode.ts

import { useState, useEffect, useRef } from 'react';
import { useChatWebSocket } from 'hooks/useChatWebSocket';
import { useSettings } from '@/state/useSettings';
import { useVisibleOpenCodeModels } from '@/state/useModels';
import { useOpenCodeSessions } from '@/state/useOpenCodeSessions';
import type { ChatMessage, ServerMessage, TaskInfo } from './types';

type UseOpenCodeOptions = {
  replaceUrl?: boolean;
  taskInfo?: TaskInfo;
};

export const useOpenCode = (initialSessionId?: string, initialModel?: string | null, options?: UseOpenCodeOptions) => {
  const { replaceUrl = true, 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 streamingRef = useRef('');
  const rafRef = useRef<number | null>(null);
  const sessionIdRef = useRef<string | null>(initialSessionId ?? null);
  const selectedModelRef = useRef<string | null>(initialModel ?? null);

  const updateSelectedModel = (value: string | null) => {
    selectedModelRef.current = value;
    setSelectedModel(value);
  };

  const { getMessages } = useOpenCodeSessions();
  const { settings } = useSettings();
  const openCodeModels = useVisibleOpenCodeModels();

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

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

  const commitStreaming = () => {
    // Cancel any pending RAF to prevent stale reads of cleared streamingRef
    if (rafRef.current !== null) {
      cancelAnimationFrame(rafRef.current);
      rafRef.current = null;
    }
    if (!streamingRef.current) return;
    setMessages((prev) => [...prev, { role: 'assistant', text: streamingRef.current }]);
    streamingRef.current = '';
    setStreamingText('');
  };

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

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

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

      case 'assistant:text':
        // Server sends the final complete text — discard streaming and use this instead
        if (rafRef.current !== null) {
          cancelAnimationFrame(rafRef.current);
          rafRef.current = null;
        }
        streamingRef.current = '';
        setStreamingText('');
        setMessages((prev) => [...prev, { role: 'assistant', text: msg.text }]);
        break;

      case 'tool:use':
        commitStreaming();
        setMessages((prev) => {
          const existing = prev.find((m) => m.role === 'tool' && m.toolUseId === msg.toolUseId);
          if (existing) {
            // Update input (running event sends actual input after pending)
            return prev.map((m) =>
              m.role === 'tool' && m.toolUseId === msg.toolUseId
                ? { ...m, toolName: msg.toolName, toolInput: msg.toolInput }
                : m,
            );
          }
          return [
            ...prev,
            { role: 'tool', toolName: msg.toolName, toolInput: msg.toolInput, toolUseId: msg.toolUseId },
          ];
        });
        break;

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

      case 'result':
        commitStreaming();
        setMessages((prev) => [
          ...prev,
          {
            role: 'result',
            costUsd: msg.costUsd,
            durationMs: msg.durationMs,
            numTurns: msg.numTurns,
            isError: msg.isError,
          },
        ]);
        setIsGenerating(false);
        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 OpenCode on mount when resuming a session
  useEffect(() => {
    if (!initialSessionId) return;
    getMessages(initialSessionId)
      .then((data) => {
        if (Array.isArray(data) && data.length > 0) setMessages(data);
      })
      .catch(() => {});
  }, [initialSessionId]);

  useEffect(() => {
    selectedModelRef.current = selectedModel;
  }, [selectedModel]);

  // Seed default model for OpenCode if none selected
  useEffect(() => {
    if (selectedModel) return;
    if (settings.chat.defaultProvider !== 'opencode' || !settings.chat.defaultModel) return;
    if (!openCodeModels.some((m) => m.id === settings.chat.defaultModel)) return;
    updateSelectedModel(settings.chat.defaultModel);
  }, [openCodeModels, selectedModel, settings.chat.defaultModel, settings.chat.defaultProvider]);

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

  const sendPrompt = (text: string, attachmentIds?: string[], images?: { filename: string; dataUrl: string }[]) => {
    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);

    const modelId = selectedModelRef.current;
    const selectedOption = modelId ? openCodeModels.find((m) => m.id === modelId) : undefined;
    const payload = {
      type: 'chat',
      prompt: text,
      sessionId: sessionIdRef.current,
      ...(modelId
        ? {
            model: {
              modelID: modelId,
              ...(selectedOption?.providerId ? { providerID: selectedOption.providerId } : {}),
            },
          }
        : {}),
      ...(attachmentIds?.length ? { attachmentIds } : {}),
      ...(imageData?.length ? { images: imageData } : {}),
      ...(taskInfo ? { taskInfo } : {}),
    };
    console.log('[opencode-ui] ws send', payload);
    send(payload);
  };

  const stopGeneration = () => {
    send({ type: 'stop' });
  };

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