public interface StreamingChatModel extends Model<Prompt, ChatResponse> {
    // 同步方法（继承自 Model）
    ChatResponse call(Prompt prompt);
    
    // 流式方法
    default Flux<ChatResponse> stream(Prompt prompt) {
        return stream(prompt, null);
    }
    
    Flux<ChatResponse> stream(Prompt prompt, ChatResponse previousChatResponse);
}

// OpenAiChatModel.java（简化版）
public class OpenAiChatModel implements ChatModel, StreamingChatModel {
    
    private final OpenAiApi api;  // 底层 HTTP 客户端
    
    @Override
    public Flux<ChatResponse> stream(Prompt prompt, ChatResponse previousChatResponse) {
        // 1. 构建请求参数
        ChatCompletionRequest request = createRequest(prompt, true);  // stream=true
        
        // 2. 调用 OpenAI API，获取 SSE 流
        Flux<ChatCompletionChunk> chunkFlux = this.api.chatCompletionStream(request);
        
        // 3. 将 chunk 流转换为 ChatResponse 流
        return chunkFlux.map(this::toChatResponse)
                       .switchOnFirst((signal, flux) -> {
                           // 处理第一个 chunk（包含 role 信息）
                           // 后续 chunk 包含 content 片段
                       });
    }
}

@RestController
@RequestMapping("/api/chat")
public class StreamChatController {

    private final StreamingChatModel streamingChatModel;

    public StreamChatController(StreamingChatModel streamingChatModel) {
        this.streamingChatModel = streamingChatModel;
    }

    /**
     * SSE 流式聊天接口
     * produces = MediaType.TEXT_EVENT_STREAM_VALUE 表示返回 SSE 流
     */
    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> streamChat(@RequestParam String message) {
        Prompt prompt = new Prompt(new UserMessage(message));

        return streamingChatModel.stream(prompt)
                .map(response -> {
                    String content = response.getResult().getOutput().getText();
                    return ServerSentEvent.<String>builder()
                            .data(content != null ? content : "")
                            .build();
                })
                .concatWith(Flux.just(
                        ServerSentEvent.<String>builder()
                                .event("done")
                                .data("[DONE]")
                                .build()
                ))
                .onErrorResume(e -> {
                    // 流式错误处理：发送错误事件后关闭
                    return Flux.just(
                            ServerSentEvent.<String>builder()
                                    .event("error")
                                    .data("流式输出异常: " + e.getMessage())
                                    .build()
                    );
                });
    }
}

// 简化的聚合逻辑
public class ChatCompletionChunkAggregator {
    
    private final StringBuilder contentBuilder = new StringBuilder();
    private final Map<Integer, ToolCallBuilder> toolCallBuilders = new HashMap<>();
    
    public ChatResponse aggregate(ChatCompletionChunk chunk) {
        // 1. 聚合文本内容
        String delta = chunk.getChoices().get(0).getDelta().getContent();
        if (delta != null) {
            contentBuilder.append(delta);
        }
        
        // 2. 聚合 Tool Call
        List<ToolCallDelta> toolCalls = chunk.getChoices().get(0).getDelta().getToolCalls();
        if (toolCalls != null) {
            for (ToolCallDelta tc : toolCalls) {
                toolCallBuilders.computeIfAbsent(tc.getIndex(), k -> new ToolCallBuilder())
                    .appendName(tc.getFunction().getName())
                    .appendArguments(tc.getFunction().getArguments());
            }
        }
        
        // 3. 当 chunk 的 finish_reason 不为 null 时，表示流结束
        if (chunk.getChoices().get(0).getFinishReason() != null) {
            return buildFinalResponse();
        }
        
        return null;  // 还没结束，继续聚合
    }
}

/**
 * 带 Tool Call 的流式输出
 * Spring AI 会自动处理 Tool Call 的聚合和执行
 */
@GetMapping(value = "/stream-with-tools", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
public Flux<ServerSentEvent<String>> streamWithTools(@RequestParam String message) {
    Prompt prompt = new Prompt(
        new UserMessage(message),
        OpenAiChatOptions.builder()
            .function("getWeather")      // 注册可用工具
            .function("searchDatabase")
            .build()
    );

    return streamingChatModel.stream(prompt)
            .map(response -> {
                String content = response.getResult().getOutput().getText();
                // 注意：Tool Call 执行期间，content 可能为 null
                // 但 Tool Call 完成后的最终响应会包含完整文本
                if (content != null && !content.isEmpty()) {
                    return ServerSentEvent.<String>builder()
                            .data(content)
                            .build();
                }
                return ServerSentEvent.<String>builder()
                        .event("tool-call")
                        .data("正在调用工具...")
                        .build();
            });
}

// LangChain4j 的流式接口
public interface StreamingChatLanguageModel extends ChatLanguageModel {
    void generate(List<ChatMessage> messages, TokenStreamHandler handler);
}

// TokenStream 的使用方式
public interface TokenStream {
    TokenStream onPartialResponse(Consumer<String> handler);  // 每个 token 的回调
    TokenStream onToolCalled(Consumer<ToolExecution> handler); // Tool Call 回调
    TokenStream onCompleteResponse(Consumer<ChatResponse> handler); // 完成回调
    TokenStream onError(ThrowableHandler handler);  // 错误回调
    void start();  // 开始生成
}

StreamingChatLanguageModel model = OpenAiStreamingChatModel.builder()
        .apiKey(System.getenv("OPENAI_API_KEY"))
        .modelName("gpt-4o")
        .build();

TokenStream tokenStream = model.generate(
    List.of(new UserMessage("你好")),
    handler -> handler
        .onPartialResponse(token -> System.out.print(token))  // 逐 token 输出
        .onCompleteResponse(response -> System.out.println("\n完成"))
        .onError(error -> error.printStackTrace())
);
tokenStream.start();

streamingChatModel.stream(prompt)
    // 策略 1：遇到错误时返回一个兜底响应
    .onErrorResume(e -> {
        log.error("流式输出异常", e);
        ChatResponse fallback = new ChatResponse(
            List.of(new Generation(new AssistantMessage("抱歉，服务暂时不可用，请稍后重试。")))
        );
        return Flux.just(fallback);
    })
    // 策略 2：超时控制
    .timeout(Duration.ofSeconds(30))
    .onErrorResume(TimeoutException.class, e -> {
        return Flux.just(new ChatResponse(
            List.of(new Generation(new AssistantMessage("响应超时，请重试。")))
        );
    });

streamingChatModel.stream(prompt)
    // 策略 1：丢弃旧数据（适合实时性要求高的场景）
    .onBackpressureDrop(overflowed -> {
        log.warn("背压丢弃: {}", overflowed);
    })
    // 策略 2：缓冲（适合不能丢数据的场景）
    .onBackpressureBuffer(1000, BufferOverflowStrategy.DROP_OLDEST)
    // 策略 3：批量发送（减少网络请求）
    .bufferTimeout(10, Duration.ofMillis(100))
    .map(tokens -> String.join("", tokens));

@Configuration
public class StreamingConfig {
    
    @Bean
    public WebClient.Builder webClientBuilder() {
        return WebClient.builder()
            .defaultHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
            .clientConnector(new ReactorClientHttpConnector(
                HttpClient.create()
                    .option(ChannelOption.CONNECT_TIMEOUT_MILLIS, 5000)
                    .responseTimeout(Duration.ofSeconds(60))
                    .maxConnections(200)  // 最大连接数
            ));
    }
}

@Service
public class RateLimitedChatService {

    private final Semaphore semaphore = new Semaphore(50);  // 最多 50 个并发
    private final StreamingChatModel streamingChatModel;

    public Flux<ChatResponse> streamWithRateLimit(Prompt prompt) {
        return Flux.defer(() -> {
            if (semaphore.tryAcquire()) {
                return streamingChatModel.stream(prompt)
                    .doFinally(signal -> semaphore.release());
            } else {
                return Flux.error(new RuntimeException("请求过于繁忙，请稍后重试"));
            }
        });
    }
}

@Component
public class StreamMetrics {

    private final MeterRegistry registry;
    private final Timer ttftTimer;
    private final Counter streamCounter;
    private final AtomicInteger activeStreams = new AtomicInteger(0);

    public StreamMetrics(MeterRegistry registry) {
        this.registry = registry;
        this.ttftTimer = Timer.builder("agent.stream.ttft")
                .description("Time to first token")
                .publishPercentiles(0.5, 0.95, 0.99)
                .register(registry);
        this.streamCounter = Counter.builder("agent.stream.total")
                .description("Total stream requests")
                .register(registry);
    }

    public void recordStreamStart() {
        activeStreams.incrementAndGet();
        streamCounter.increment();
    }

    public void recordTTFT(Duration duration) {
        ttftTimer.record(duration);
    }

    public void recordStreamEnd(boolean completed) {
        activeStreams.decrementAndGet();
        Counter.builder("agent.stream.completion")
                .tag("status", completed ? "success" : "interrupted")
                .register(registry)
                .increment();
    }
}

@RestController
@RequestMapping("/api/agent")
public class StreamingAgentController {

    private final StreamingChatModel chatModel;
    private final ChatMemory chatMemory;
    private final WeatherService weatherService;
    private final StreamMetrics streamMetrics;

    public StreamingAgentController(
            StreamingChatModel chatModel,
            ChatMemory chatMemory,
            WeatherService weatherService,
            StreamMetrics streamMetrics) {
        this.chatModel = chatModel;
        this.chatMemory = chatMemory;
        this.weatherService = weatherService;
        this.streamMetrics = streamMetrics;
    }

    /**
     * 流式 Agent 对话接口
     * 支持记忆、工具调用、流式输出
     */
    @GetMapping(value = "/chat/{conversationId}", 
                produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> streamAgentChat(
            @PathVariable String conversationId,
            @RequestParam String message) {
        
        streamMetrics.recordStreamStart();
        Instant startTime = Instant.now();

        // 1. 从记忆中获取历史消息
        List<Message> history = chatMemory.get(conversationId);
        
        // 2. 构建包含历史和工具的 Prompt
        List<Message> messages = new ArrayList<>(history);
        messages.add(new UserMessage(message));
        
        Prompt prompt = new Prompt(messages,
            OpenAiChatOptions.builder()
                .function("getWeather")
                .function("searchKnowledge")
                .build()
        );

        // 3. 流式调用
        return chatModel.stream(prompt)
                .index()  // 添加索引，用于判断是否是第一个 token
                .map(tuple -> {
                    long index = tuple.getT1();
                    ChatResponse response = tuple.getT2();
                    
                    // 记录 TTFT
                    if (index == 0) {
                        streamMetrics.recordTTFT(
                            Duration.between(startTime, Instant.now())
                        );
                    }
                    
                    String content = response.getResult().getOutput().getText();
                    return ServerSentEvent.<String>builder()
                            .id(String.valueOf(index))
                            .data(content != null ? content : "")
                            .build();
                })
                // 发送完成信号
                .concatWith(Flux.just(
                    ServerSentEvent.<String>builder()
                            .event("done")
                            .data("[DONE]")
                            .build()
                ))
                // 保存对话到记忆
                .doOnComplete(() -> {
                    chatMemory.add(conversationId, new UserMessage(message));
                    // 注意：流式场景下，需要在流完成后保存完整的 assistant 回复
                    // 这里简化处理，实际需要聚合完整响应
                    streamMetrics.recordStreamEnd(true);
                })
                // 错误处理
                .onErrorResume(e -> {
                    log.error("流式 Agent 对话异常, conversationId={}", conversationId, e);
                    streamMetrics.recordStreamEnd(false);
                    return Flux.just(
                        ServerSentEvent.<String>builder()
                                .event("error")
                                .data("服务异常: " + e.getMessage())
                                .build()
                    );
                });
    }
}

// 聚合完整响应后再保存到记忆
AtomicReference<StringBuilder> fullResponse = new AtomicReference<>(new StringBuilder());

chatModel.stream(prompt)
    .doOnNext(response -> {
        String content = response.getResult().getOutput().getText();
        if (content != null) {
            fullResponse.get().append(content);
        }
    })
    .doOnComplete(() -> {
        // 流结束，保存完整对话
        chatMemory.add(conversationId, new UserMessage(message));
        chatMemory.add(conversationId, 
            new AssistantMessage(fullResponse.get().toString()));
    });