beige/frontend/src/hooks/useSse.ts

'use client';
/**
 * useSse — SSE 客户端 React hook
 * 接收 SSE 事件，分派到 generationStore
 * event_seq 去重 + 断线重连UI
 * 五维分：SSE 推 total，五维详情通过 GET /tasks/{id} 补全
 */
import { useEffect, useRef } from 'react';
import { createSseClient, SseClient, SseConnectionStatus } from '@/lib/sse';
import { SseEvent } from '@/types/sse';
import { useGenerationStore } from '@/stores/generationStore';
import { TextCandidate, ImageCandidate } from '@/types/index';
import { api } from '@/lib/api';

export function useSse(taskId: number | null) {
  const clientRef = useRef<SseClient | null>(null);
  const store = useGenerationStore();

  useEffect(() => {
    if (!taskId) return;
    store.reset();

    // reset 之后立即回填已有候选（历史/已完成任务 SSE 不重放）。
    // 必须放在 reset 之后、与 SSE 同一 effect 内，否则与页面侧 hydrate 形成
    // reset 把 hydrate 数据清空的竞态（SPA 跳转进来会看到空白）。
    let cancelled = false;
    api.get<{
      text_candidates?: TextCandidate[];
      image_candidates?: ImageCandidate[];
    }>(`/api/v1/tasks/${taskId}`).then((task) => {
      if (cancelled) return;
      const texts = task.text_candidates ?? [];
      const images = task.image_candidates ?? [];
      if (texts.length || images.length) {
        store.hydrateFromTask({ textCandidates: texts, imageCandidates: images });
      }
    }).catch(() => { /* 进行中任务靠 SSE 填充 */ });

    const client = createSseClient(taskId, {
      onEvent: (event: SseEvent) => handleEvent(event, store, taskId),
      onStatusChange: (status: SseConnectionStatus, attempt?: number) => {
        store.setConnectionStatus(status, attempt);
      },
    });
    clientRef.current = client;

    return () => {
      cancelled = true;
      client.close();
      clientRef.current = null;
    };
  }, [taskId]);

  return {
    connectionStatus: store.connectionStatus,
    reconnectAttempt: store.reconnectAttempt,
  };
}

function handleEvent(
  event: SseEvent,
  store: ReturnType<typeof useGenerationStore.getState>,
  taskId: number,
) {
  switch (event.type) {
    case 'task_started':
      store.setTextProgress(0, event.data.total_text);
      store.setImageProgress(0, event.data.total_image);
      break;

    case 'text_progress':
      store.setTextProgress(event.data.done, event.data.total);
      break;

    case 'text_candidate': {
      // SSE 只推 total 分；先用 total 占位，再异步补全五维分
      const totalScore = typeof event.data.score === 'number' ? event.data.score : 0;
      const tc: TextCandidate = {
        candidate_id: event.data.candidate_id,
        angle_label: event.data.angle_label,
        content: event.data.content,
        source: 'ai',
        score: { title: 0, emotion: 0, selling: 0, keyword: 0, compliance: 0, total: totalScore },
        banned_word_status: 'pass',
      };
      store.addTextCandidate(tc);
      // 异步补全五维分（不阻塞 SSE）
      api.get<{ text_candidates?: Array<{ candidate_id: number; score?: TextCandidate['score'] }> }>(
        `/api/v1/tasks/${taskId}`
      ).then((task) => {
        const found = task.text_candidates?.find((c) => c.candidate_id === tc.candidate_id);
        if (found?.score) store.patchTextCandidateScore(tc.candidate_id, found.score);
      }).catch(() => { /* 五维分补全失败不影响主流程 */ });
      break;
    }

    case 'image_progress':
      store.setImageProgress(event.data.done, event.data.total);
      break;

    case 'image_candidate': {
      const ic: ImageCandidate = {
        candidate_id: event.data.candidate_id,
        strategy: event.data.strategy,
        url: event.data.url,
        role: event.data.role as ImageCandidate['role'],
      };
      store.addImageCandidate(ic);
      break;
    }

    case 'batch_failed':
      store.markBatchFailed(event.data.batch, event.data.reason, event.data.retryable);
      break;

    case 'flywheel_injected':
      store.setFlywheelInjected({
        recentPreference: event.data.recent_preference,
        rejectReasons: event.data.reject_reasons,
      });
      break;

    case 'error':
    case 'heartbeat':
    case 'analyze_done':
    case 'task_done':
      break;
  }
}