writer-work-flow/apps/web/lib/workbench/useChapterRewrite.ts

"use client";

import { useCallback, useReducer, useRef, useState } from "react";

import { API_BASE_PUBLIC } from "@/lib/api/config";
import {
  SseFrameBuffer,
  initialStreamState,
  reduceStream,
  type StreamState,
} from "@/lib/stream/sse";

// 一轮整章重写的产出（作者意见 + 产出正文）。多轮迭代累积成版本栈。
export interface RewriteVersion {
  feedback: string;
  text: string;
}

type Action =
  | { type: "start" }
  | { type: "events"; events: ReturnType<SseFrameBuffer["push"]> }
  | { type: "abort" }
  | { type: "fail"; code: string; message: string }
  | { type: "reset" };

function reducer(state: StreamState, action: Action): StreamState {
  switch (action.type) {
    case "start":
      return { phase: "streaming", text: "", error: null };
    case "events":
      return action.events.reduce(reduceStream, state);
    case "abort":
      return { ...state, phase: "aborted" };
    case "fail":
      return {
        ...state,
        phase: "error",
        error: { code: action.code, message: action.message },
      };
    case "reset":
      return initialStreamState;
    default:
      return state;
  }
}

export interface UseChapterRewrite {
  // 当前流式状态（text = 正在生成/最新一版正文，供 live 展示）。
  state: StreamState;
  isStreaming: boolean;
  // 已完成的版本栈（每轮一版）。
  versions: RewriteVersion[];
  latest: RewriteVersion | null;
  // 发起一轮整章重写：作者意见 + 当前整章草稿 → 流式重写一版。返回本版（done 且有正文），否则 null。
  send: (
    projectId: string,
    chapterNo: number,
    feedback: string,
    priorDraft: string,
  ) => Promise<RewriteVersion | null>;
  stop: () => void;
  reset: () => void;
}

// 整章再沟通/重写流：POST .../rewrite（feedback + prior_draft）→ 消费 SSE → 累积成新一版。
// 每轮完成压入版本栈；下一轮以最新版正文为 prior_draft + 新意见迭代（HITL：接受才落草稿）。
// 用 fetch+ReadableStream（EventSource 不支持 POST），复用 lib/stream/sse 的帧解析与归一。
export function useChapterRewrite(): UseChapterRewrite {
  const [state, dispatch] = useReducer(reducer, initialStreamState);
  const [versions, setVersions] = useState<RewriteVersion[]>([]);
  const controllerRef = useRef<AbortController | null>(null);
  const textRef = useRef("");

  const stop = useCallback(() => {
    controllerRef.current?.abort();
    controllerRef.current = null;
    dispatch({ type: "abort" });
  }, []);

  const reset = useCallback(() => {
    dispatch({ type: "reset" });
    setVersions([]);
  }, []);

  const send = useCallback<UseChapterRewrite["send"]>(
    async (projectId, chapterNo, feedback, priorDraft) => {
      const controller = new AbortController();
      controllerRef.current = controller;
      textRef.current = "";
      dispatch({ type: "start" });
      try {
        const res = await fetch(
          `${API_BASE_PUBLIC}/projects/${projectId}/chapters/${chapterNo}/rewrite`,
          {
            method: "POST",
            headers: {
              Accept: "text/event-stream",
              "Content-Type": "application/json",
            },
            body: JSON.stringify({ feedback, prior_draft: priorDraft }),
            signal: controller.signal,
          },
        );
        if (!res.ok || !res.body) {
          // 流前错误（如无凭据 → 503 LLM_UNAVAILABLE，JSON 信封而非帧）。
          let code = "STREAM_FAILED";
          let message = `整章重写请求失败（${res.status}）`;
          try {
            const body = (await res.json()) as {
              error?: { code?: string; message?: string };
            };
            if (body.error?.code) code = body.error.code;
            if (body.error?.message) message = body.error.message;
          } catch {
            // 非 JSON 信封，沿用默认文案。
          }
          dispatch({ type: "fail", code, message });
          return null;
        }
        const reader = res.body.getReader();
        const decoder = new TextDecoder();
        const buffer = new SseFrameBuffer();
        let done = false;
        for (;;) {
          const { value, done: readerDone } = await reader.read();
          if (readerDone) break;
          const events = buffer.push(decoder.decode(value, { stream: true }));
          if (events.length > 0) {
            dispatch({ type: "events", events });
            for (const ev of events) {
              if (ev.event === "token") textRef.current += ev.data.text;
              if (ev.event === "done") done = true;
            }
          }
        }
        // 收到 done 且有正文 → 压入版本栈（未 done/空文本不入栈，避免脏版本）。
        if (done && textRef.current.length > 0) {
          const version: RewriteVersion = { feedback, text: textRef.current };
          setVersions((prev) => [...prev, version]);
          return version;
        }
        return null;
      } catch (err: unknown) {
        if (err instanceof DOMException && err.name === "AbortError") {
          return null; // 用户主动停止：state 已置 aborted。
        }
        const message = err instanceof Error ? err.message : "未知网络错误";
        dispatch({ type: "fail", code: "NETWORK", message });
        return null;
      } finally {
        controllerRef.current = null;
      }
    },
    [],
  );

  return {
    state,
    isStreaming: state.phase === "streaming",
    versions,
    latest: versions.at(-1) ?? null,
    send,
    stop,
    reset,
  };
}