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

import { streamGroundedChat } from "../api";
import { ChatStreamError } from "../sse";
import {
  INITIAL_CHAT_STATE,
  type ChatCompletionRequest,
  type ChatRunState,
  type ChatStreamEvent,
} from "../types";

function stateForEvent(current: ChatRunState, event: ChatStreamEvent): ChatRunState {
  switch (event.name) {
    case "meta":
      return { ...current, meta: event };
    case "retrieval":
      return { ...current, phase: "generating", retrieval: event };
    case "delta":
      return { ...current, phase: "generating", answer: current.answer + event.text };
    case "citations":
      return { ...current, citations: event.citations };
    case "usage":
      return { ...current, usage: event };
    case "done":
      return {
        ...current,
        phase:
          event.answer_mode === "grounded"
            ? "complete"
            : event.answer_mode === "refused"
              ? "refused"
              : "retrieval_only",
        done: event,
      };
    case "error":
      return {
        ...current,
        phase: "error",
        citations: current.retrieval?.evidence ?? [],
        streamError: event,
        errorMessage: event.retryable
          ? "生成模型暂不可用，已保留本次检索证据，可以稍后重试。"
          : "生成过程未能安全完成，已切换为仅检索证据模式。",
      };
  }
}

export function useGroundedChat() {
  const [state, setState] = useState<ChatRunState>(INITIAL_CHAT_STATE);
  const controllerRef = useRef<AbortController | null>(null);
  const runIdRef = useRef(0);

  const stop = useCallback(() => {
    controllerRef.current?.abort();
  }, []);

  const start = useCallback(async (request: ChatCompletionRequest) => {
    controllerRef.current?.abort();
    const controller = new AbortController();
    controllerRef.current = controller;
    const runId = ++runIdRef.current;
    setState({ ...INITIAL_CHAT_STATE, phase: "retrieving", request });

    try {
      await streamGroundedChat(request, {
        signal: controller.signal,
        onEvent: (event) => {
          if (runId === runIdRef.current) setState((current) => stateForEvent(current, event));
        },
      });
    } catch (error) {
      if (runId !== runIdRef.current) return;
      if (error instanceof ChatStreamError && error.kind === "aborted") {
        setState((current) => ({
          ...current,
          phase: "stopped",
          errorMessage: null,
        }));
      } else {
        setState((current) => ({
          ...current,
          phase: "error",
          answer:
            error instanceof ChatStreamError && error.kind === "invalid_stream"
              ? ""
              : current.answer,
          citations:
            error instanceof ChatStreamError && error.kind === "invalid_stream"
              ? []
              : current.citations.length > 0
                ? current.citations
                : (current.retrieval?.evidence ?? []),
          errorMessage:
            error instanceof ChatStreamError
              ? error.message
              : "回答流发生未预期错误，请重新发起问题。",
        }));
      }
    } finally {
      if (runId === runIdRef.current) controllerRef.current = null;
    }
  }, []);

  useEffect(() => () => controllerRef.current?.abort(), []);

  return {
    state,
    isRunning: state.phase === "retrieving" || state.phase === "generating",
    start,
    stop,
  };
}