import dev.langchain4j.model.chat.listener.ChatModelErrorContext;
import dev.langchain4j.model.chat.listener.ChatModelListener;
import dev.langchain4j.model.chat.listener.ChatModelRequestContext;
import dev.langchain4j.model.chat.listener.ChatModelResponseContext;

@Component
public class AgentObservabilityListener implements ChatModelListener {

    private static final Logger log = LoggerFactory.getLogger(AgentObservabilityListener.class);
    private final MeterRegistry meterRegistry;

    public AgentObservabilityListener(MeterRegistry meterRegistry) {
        this.meterRegistry = meterRegistry;
    }

    @Override
    public void onRequest(ChatModelRequestContext context) {
        // 请求发出前：记录开始时间，注入 Trace ID
        context.attributes().put("startTime", System.nanoTime());
        String traceId = MDC.get("traceId");
        if (traceId != null) {
            context.attributes().put("traceId", traceId);
        }
        log.debug("LLM request initiated, model: {}", context.chatRequest().modelName());
    }

    @Override
    public void onResponse(ChatModelResponseContext context) {
        // 响应返回后：记录 Token 消耗、延迟、成本
        long startTime = (long) context.attributes().getOrDefault("startTime", 0L);
        long latencyMs = (System.nanoTime() - startTime) / 1_000_000;

        var tokenUsage = context.chatResponse().tokenUsage();
        String model = context.chatResponse().modelName();

        // 记录指标
        meterRegistry.counter("llm.requests.total",
                "model", model,
                "status", "success"
        ).increment();

        meterRegistry.counter("llm.tokens.total",
                "model", model,
                "type", "prompt"
        ).increment(tokenUsage.inputTokenCount());

        meterRegistry.counter("llm.tokens.total",
                "model", model,
                "type", "completion"
        ).increment(tokenUsage.outputTokenCount());

        meterRegistry.timer("llm.latency",
                "model", model
        ).record(Duration.ofMillis(latencyMs));

        // 计算并记录成本
        double cost = calculateCost(model,
                tokenUsage.inputTokenCount(),
                tokenUsage.outputTokenCount());
        meterRegistry.counter("llm.cost.total",
                "model", model
        ).increment(cost);

        log.info("LLM call completed: model={}, promptTokens={}, completionTokens={}, " +
                "latencyMs={}, cost=${}",
                model,
                tokenUsage.inputTokenCount(),
                tokenUsage.outputTokenCount(),
                latencyMs,
                String.format("%.4f", cost));
    }

    @Override
    public void onError(ChatModelErrorContext context) {
        meterRegistry.counter("llm.requests.total",
                "model", "unknown",
                "status", "error"
        ).increment();

        log.error("LLM call failed: {}", context.error().getMessage());
    }

    private double calculateCost(String model, long promptTokens, long completionTokens) {
        // 2026 年主流模型定价（每 1M tokens）
        return switch (model) {
            case "gpt-4o" -> promptTokens * 2.5 / 1_000_000 + completionTokens * 10.0 / 1_000_000;
            case "gpt-4o-mini" -> promptTokens * 0.15 / 1_000_000 + completionTokens * 0.6 / 1_000_000;
            case "claude-3.5-sonnet" -> promptTokens * 3.0 / 1_000_000 + completionTokens * 15.0 / 1_000_000;
            default -> 0.0;
        };
    }
}

// 方式一：编程式——构建 ChatModel 时注入
@Bean
public ChatLanguageModel chatModel(AgentObservabilityListener listener) {
    return OpenAiChatModel.builder()
            .apiKey(System.getenv("OPENAI_API_KEY"))
            .modelName("gpt-4o")
            .listeners(List.of(listener))
            .build();
}

// 方式二：声明式（AiService）——通过 builder 注入
AiService.builder(CustomerServiceAgent.class)
        .chatLanguageModel(chatModel)
        .chatMemoryProvider(memoryProvider)
        .tools(orderTools)
        .build();

import dev.langchain4j.opentelemetry.OpenTelemetryChatModelListener;
import io.opentelemetry.api.OpenTelemetry;
import io.opentelemetry.api.trace.Tracer;

@Configuration
public class ObservabilityConfig {

    @Bean
    public OpenTelemetryChatModelListener otelListener(OpenTelemetry openTelemetry) {
        Tracer tracer = openTelemetry.getTracer("langchain4j-agent");
        return new OpenTelemetryChatModelListener(tracer);
    }
}

import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.Tracer;

@Component
public class OrderTools {

    private final OrderService orderService;
    private final Tracer tracer;

    public OrderTools(OrderService orderService, Tracer tracer) {
        this.orderService = orderService;
        this.tracer = tracer;
    }

    @Tool("根据订单号查询订单状态")
    public OrderStatus getOrderStatus(String orderId) {
        Span currentSpan = Span.current();
        currentSpan.setAttribute("order.id", orderId);

        long start = System.nanoTime();
        try {
            OrderStatus status = orderService.queryStatus(orderId);
            currentSpan.setAttribute("order.status", status.getStatus());
            return status;
        } catch (Exception e) {
            currentSpan.setAttribute("error.type", e.getClass().getSimpleName());
            throw e;
        } finally {
            long duration = (System.nanoTime() - start) / 1_000_000;
            currentSpan.setAttribute("tool.duration_ms", duration);
        }
    }
}

import io.micrometer.observation.Observation;
import io.micrometer.observation.ObservationRegistry;

@Component
public class AgentInteractionObserver {

    private final ObservationRegistry observationRegistry;

    public AgentInteractionObserver(ObservationRegistry observationRegistry) {
        this.observationRegistry = observationRegistry;
    }

    public <T> T observeAgentInteraction(String userId, String agentType,
                                          Supplier<T> interaction) {
        return Observation.createNotStarted("agent.interaction", observationRegistry)
                .contextualName("agent-" + agentType)
                .lowCardinalityKeyValue("agent.type", agentType)
                .highCardinalityKeyValue("user.id", userId)
                .observe(() -> {
                    long startTime = System.nanoTime();
                    try {
                        T result = interaction.get();
                        long duration = (System.nanoTime() - startTime) / 1_000_000;
                        Observation observation = Observation.start("agent.interaction", observationRegistry);
                        observation.highCardinalityKeyValue("duration_ms", String.valueOf(duration));
                        observation.stop();
                        return result;
                    } catch (Exception e) {
                        throw e;
                    }
                });
    }
}

@Service
public class CustomerServiceAgent {

    private final ChatClient chatClient;
    private final AgentInteractionObserver observer;

    public String handleInquiry(String userId, String message) {
        return observer.observeAgentInteraction(userId, "customer-service", () -> {
            return chatClient.prompt()
                    .system("你是客服 Agent，帮助用户处理订单问题")
                    .user(message)
                    .call()
                    .content();
        });
    }
}

@Component
public class CostAlertService {

    private final MeterRegistry meterRegistry;
    private final NotificationService notificationService;

    // 告警阈值
    private static final double HOURLY_COST_LIMIT = 50.0;   // 单小时 $50
    private static final double DAILY_COST_LIMIT = 500.0;    // 单日 $500
    private static final long TOKENS_PER_REQUEST_LIMIT = 10000; // 单次请求 Token 上限

    @Scheduled(fixedRate = 60_000) // 每分钟检查
    public void checkCostAlerts() {
        // 检查小时级成本
        double hourlyCost = getMetricValue("llm.cost.total", Duration.ofHours(1));
        if (hourlyCost > HOURLY_COST_LIMIT) {
            notificationService.sendAlert(
                AlertLevel.CRITICAL,
                String.format("⚠️ AI Agent 小时成本超限: $%.2f (阈值: $%.2f)", hourlyCost, HOURLY_COST_LIMIT)
            );
        }

        // 检查日级成本
        double dailyCost = getMetricValue("llm.cost.total", Duration.ofDays(1));
        if (dailyCost > DAILY_COST_LIMIT) {
            notificationService.sendAlert(
                AlertLevel.CRITICAL,
                String.format("🚨 AI Agent 日成本超限: $%.2f (阈值: $%.2f)", dailyCost, DAILY_COST_LIMIT)
            );
        }
    }

    @Scheduled(fixedRate = 10_000) // 每 10 秒检查（实时检测异常请求）
    public void checkAnomalousRequests() {
        // 检测异常大的单次请求（可能是推理死循环）
        double maxTokens = meterRegistry.get("llm.tokens.per.request")
                .tag("statistic", "MAX")
                .gauge().value();
        if (maxTokens > TOKENS_PER_REQUEST_LIMIT) {
            notificationService.sendAlert(
                AlertLevel.WARNING,
                String.format("⚠️ 检测到异常大的请求: %d tokens (阈值: %d)", (long) maxTokens, TOKENS_PER_REQUEST_LIMIT)
            );
        }
    }

    private double getMetricValue(String metricName, Duration duration) {
        // 从 MeterRegistry 获取时间窗口内的累计值
        // 具体实现取决于你使用的 Registry 类型
        return meterRegistry.get(metricName).counter().count();
    }
}

// 简化版源码分析
public ChatResponse generate(ChatRequest request) {
    ChatModelRequestContext requestContext = new ChatModelRequestContext(request, attributes);
    
    // 1. 触发所有 Listener 的 onRequest
    for (ChatModelListener listener : listeners) {
        try {
            listener.onRequest(requestContext);
        } catch (Exception e) {
            log.warn("Error in ChatModelListener.onRequest", e);
        }
    }
    
    try {
        // 2. 执行实际的 LLM 调用
        ChatResponse response = doGenerate(request);
        
        ChatModelResponseContext responseContext = new ChatModelResponseContext(
            request, response, attributes);
        
        // 3. 触发所有 Listener 的 onResponse
        for (ChatModelListener listener : listeners) {
            try {
                listener.onResponse(responseContext);
            } catch (Exception e) {
                log.warn("Error in ChatModelListener.onResponse", e);
            }
        }
        
        return response;
    } catch (Exception e) {
        ChatModelErrorContext errorContext = new ChatModelErrorContext(e, attributes);
        
        // 4. 出错时触发 onError
        for (ChatModelListener listener : listeners) {
            try {
                listener.onError(errorContext);
            } catch (Exception ex) {
                log.warn("Error in ChatModelListener.onError", ex);
            }
        }
        throw e;
    }
}

@Configuration
@EnableScheduling
public class AgentObservabilityConfig {

    @Bean
    public ChatModelListener chatModelListener(MeterRegistry meterRegistry) {
        return new AgentObservabilityListener(meterRegistry);
    }

    @Bean
    public ChatLanguageModel chatModel(ChatModelListener listener) {
        return OpenAiChatModel.builder()
                .apiKey(System.getenv("OPENAI_API_KEY"))
                .modelName("gpt-4o")
                .listeners(List.of(listener))
                .build();
    }

    @Bean
    public MeterFilter costTagFilter() {
        // 为所有 LLM 相关指标添加应用标签
        return MeterFilter.commonTags(Tags.of("component", "ai-agent"));
    }
}

@Component
public class AgentObservabilityListener implements ChatModelListener {

    private final MeterRegistry meterRegistry;
    private final ExecutorService asyncWriter = Executors.newFixedThreadPool(2);

    @Override
    public void onResponse(ChatModelResponseContext context) {
        // 快速：更新内存中的指标
        meterRegistry.counter("llm.tokens.total").increment(tokenUsage.inputTokenCount());

        // 异步：持久化详细日志
        asyncWriter.submit(() -> {
            llmCallLogRepository.save(buildLogFromContext(context));
        });
    }
}

@Bean
public Sampler traceSampler() {
    // 错误 100% 采样，正常 10% 采样
    return Sampler.parentBased(
        Sampler.traceIdRatioBased(0.1)  // 基础采样率 10%
    );
}

@Scheduled(cron = "0 0 2 * * *") // 每天凌晨 2 点
public void calibrateCost() {
    double providerCost = openAIBillingApi.getDailyCost(LocalDate.now().minusDays(1));
    double monitoredCost = getMonitoredCost(LocalDate.now().minusDays(1));
    double discrepancy = Math.abs(providerCost - monitoredCost) / providerCost;
    
    if (discrepancy > 0.1) { // 偏差超过 10%
        log.warn("Cost calibration off by {}%, provider: ${}, monitored: ${}",
                String.format("%.1f", discrepancy * 100),
                String.format("%.2f", providerCost),
                String.format("%.2f", monitoredCost));
    }
}