import React, { useState, useRef } from 'react';

// 定義消息結(jié)構(gòu)
interface Message {
  id: string;
  role: 'user' | 'assistant';
  content: string;
  status?: 'complete' | 'stopped'; // 標記消息狀態(tài)
}

const ChatComponent: React.FC = () => {
  const [messages, setMessages] = useState<Message[]>([]);
  const [input, setInput] = useState('');
  // 用于中斷請求的控制器
  const abortControllerRef = useRef<AbortController | null>(null);
  const [isGenerating, setIsGenerating] = useState(false);

  // 發(fā)送消息的核心邏輯
  const handleSend = async () => {
    if (!input.trim() || isGenerating) return;

    const userMessage: Message = { id: Date.now().toString(), role: 'user', content: input };
    // 初始化AI回復(fù)占位符
    const aiMessage: Message = { id: (Date.now() + 1).toString(), role: 'assistant', content: '', status: 'complete' };

    // 更新本地狀態(tài)
    setMessages(prev => [...prev, userMessage, aiMessage]);
    setInput('');
    setIsGenerating(true);

    // 創(chuàng)建新的 AbortController
    abortControllerRef.current = new AbortController();

    try {
      // 模擬流式請求 (實際項目中替換為 fetch SSE 接口)
      const response = await fetch('/api/chat', {
        method: 'POST',
        body: JSON.stringify({ prompt: input }),
        signal: abortControllerRef.current.signal, // 關(guān)鍵：綁定中斷信號
        headers: { 'Content-Type': 'application/json' }
      });

      // 模擬流式讀取
      const reader = response.body?.getReader();
      // ... 此處省略具體的流式解碼邏輯，重點在于數(shù)據(jù)追加 ...

    } catch (error: any) {
      if (error.name === 'AbortError') {
        console.log('用戶主動停止了生成');
        // 更新最后一條消息的狀態(tài)為 stopped
        setMessages(prev => {
          const newMsgs = [...prev];
          const lastMsg = newMsgs[newMsgs.length - 1];
          if (lastMsg.role === 'assistant') lastMsg.status = 'stopped';
          return newMsgs;
        });
      }
    } finally {
      setIsGenerating(false);
      abortControllerRef.current = null;
    }
  };

  // 【核心功能1】停止生成
  const handleStop = () => {
    if (abortControllerRef.current) {
      abortControllerRef.current.abort(); // 觸發(fā)中斷
    }
  };

  // 【核心功能2】重新回答
  const handleRetry = async (retryIndex: number) => {
    // 1. 找到對應(yīng)的用戶提問 (當(dāng)前索引-1)
    const userPrompt = messages[retryIndex - 1].content;

    // 2. 狀態(tài)回滾：移除當(dāng)前AI回答及之后的對話 (防止上下文污染)
    // 注意：這里保留用戶提問，只移除AI回答，也可以根據(jù)業(yè)務(wù)需求移除兩者
    setMessages(prev => prev.slice(0, retryIndex));

    // 3. 重新觸發(fā)發(fā)送邏輯 (這里簡化處理，實際應(yīng)復(fù)用 handleSend 邏輯)
    // 實際開發(fā)中建議封裝一個 receiveStream(prompt) 方法供調(diào)用
    console.log(`Retrying with prompt: ${userPrompt}`);
    // await handleSendInternal(userPrompt); 
  };

  return (
    <div className="chat-container">
      {/* 渲染消息列表 */}
      {messages.map((msg, index) => (
        <div key={msg.id} className={`message ${msg.role}`}>
          {msg.content}
          {msg.status === 'stopped' && <span className="tag"> (已停止)</span>}
          {msg.role === 'assistant' && (
             <button onClick={() => handleRetry(index)}>重新生成</button>
          )}
        </div>
      ))}

      {/* 底部操作欄 */}
      <div className="input-area">
        {isGenerating ? (
          <button onClick={handleStop}>停止生成</button>
        ) : (
          <button onClick={handleSend}>發(fā)送</button>
        )}
      </div>
    </div>
  );
};

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import asyncio

app = FastAPI()

async def generate_stream(prompt: str, request: Request):
    # 模擬LLM生成的token流
    words = ["你好", "，", "我是", "AI", "助手", "。"]
    for word in words:
        # 【關(guān)鍵點】檢測客戶端是否斷開連接
        if await request.is_disconnected():
            print("客戶端已斷開，停止生成")
            break

        # 模擬生成延遲
        await asyncio.sleep(0.5) 
        yield f"data: {word}\n\n"

@app.post("/api/chat")
async def chat(request: Request):
    data = await request.json()
    prompt = data.get("prompt")
    return StreamingResponse(
        generate_stream(prompt, request), 
        media_type="text/event-stream"
    )