package com.superbiz.agent.hook;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import reactor.core.publisher.Flux;

/**
 * ChatModel 包装器 — 捕获每次模型调用的实际 token 用量
 * 通过 TokenUsageHolder 传递给 AgentLoggingHook
 */
public class TokenTrackingChatModel implements ChatModel {

    private static final Logger log = LoggerFactory.getLogger(TokenTrackingChatModel.class);

    private final ChatModel delegate;

    public TokenTrackingChatModel(ChatModel delegate) {
        this.delegate = delegate;
    }

    @Override
    public ChatResponse call(Prompt prompt) {
        ChatResponse response = delegate.call(prompt);
        captureTokenUsage(response);
        return response;
    }

    @Override
    public Flux<ChatResponse> stream(Prompt prompt) {
        return delegate.stream(prompt);
    }

    private void captureTokenUsage(ChatResponse response) {
        try {
            if (response.getMetadata() == null || response.getMetadata().getUsage() == null) {
                return;
            }
            var usage = response.getMetadata().getUsage();
            Integer total = usage.getTotalTokens();
            if (total != null && total > 0) {
                TokenUsageHolder.set(total);
            }
        } catch (Exception e) {
            log.debug("捕获 token 用量失败", e);
        }
    }
}