package com.infoepoch.pms.agent.observability.ai;

import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.infoepoch.pms.agent.properties.AiObservabilityProperties;
import org.slf4j.MDC;
import org.springframework.util.StringUtils;

import java.util.LinkedHashMap;
import java.util.Map;
import java.util.UUID;
import java.util.function.Supplier;

/**
 * AI 观测公共能力。
 */
public class AiObservationSupport {

    private static final int PREVIEW_LIMIT = 160;

    public static final String MDC_SERVICE_NAME = "service.name";
    public static final String MDC_TRACE_ID = "trace.id";
    public static final String MDC_REQUEST_ID = "request.id";
    public static final String MDC_SESSION_ID = "session.id";
    public static final String MDC_FLOW_ID = "ai.flow.id";
    public static final String MDC_AI_SCENE = "ai.scene";
    public static final String MDC_COMPONENT_TYPE = "ai.component.type";
    public static final String MDC_COMPONENT_NAME = "ai.component.name";
    public static final String MDC_AGENT_NAME = "ai.agent.name";
    public static final String MDC_MODEL_PROVIDER = "ai.model.provider";
    public static final String MDC_MODEL_NAME = "ai.model.name";
    public static final String MDC_TOOL_NAME = "ai.tool.name";

    private final AiObservabilityProperties properties;
    private final String serviceName;

    public AiObservationSupport(AiObservabilityProperties properties, String serviceName) {
        this.properties = properties;
        this.serviceName = serviceName;
    }

    public <T> T withObservationContext(AiObservationContext context, Supplier<T> supplier) {
        Map<String, String> previous = MDC.getCopyOfContextMap();
        try {
            applyContext(context);
            return supplier.get();
        } finally {
            restore(previous);
        }
    }

    public String renderPayload(String payload) {
        if (properties.getPayloadMode() == AiObservabilityProperties.PayloadMode.OFF) {
            return "off";
        }
        if (payload == null) {
            return "长度=0 摘要=";
        }
        if (properties.getPayloadMode() == AiObservabilityProperties.PayloadMode.FULL) {
            return payload;
        }
        String normalized = payload.replaceAll("\\s+", " ").trim();
        if (normalized.length() > PREVIEW_LIMIT) {
            normalized = normalized.substring(0, PREVIEW_LIMIT) + "...";
        }
        return "长度=%s 摘要=%s".formatted(payload.length(), normalized);
    }

    public AiObservationContext currentContext() {
        return AiObservationContext.builder()
                .traceId(MDC.get(MDC_TRACE_ID))
                .requestId(MDC.get(MDC_REQUEST_ID))
                .sessionId(MDC.get(MDC_SESSION_ID))
                .flowId(MDC.get(MDC_FLOW_ID))
                .scene(MDC.get(MDC_AI_SCENE))
                .componentType(resolveComponentType(MDC.get(MDC_COMPONENT_TYPE)))
                .componentName(MDC.get(MDC_COMPONENT_NAME))
                .agentName(MDC.get(MDC_AGENT_NAME))
                .modelProvider(MDC.get(MDC_MODEL_PROVIDER))
                .modelName(MDC.get(MDC_MODEL_NAME))
                .toolName(MDC.get(MDC_TOOL_NAME))
                .build();
    }

    public Map<String, Object> toMetadata(AiObservationContext context) {
        Map<String, Object> metadata = new LinkedHashMap<>();
        put(metadata, "traceId", context.traceId());
        put(metadata, "requestId", context.requestId());
        put(metadata, "sessionId", context.sessionId());
        put(metadata, "flowId", context.flowId());
        put(metadata, "scene", context.scene());
        put(metadata, "componentType", context.componentType() == null ? null : context.componentType().name());
        put(metadata, "componentName", context.componentName());
        put(metadata, "agentName", context.agentName());
        put(metadata, "modelProvider", context.modelProvider());
        put(metadata, "modelName", context.modelName());
        put(metadata, "toolName", context.toolName());
        return metadata;
    }

    public RunnableConfig buildRunnableConfig(String threadId, AiObservationContext context) {
        RunnableConfig.Builder builder = RunnableConfig.builder();
        if (StringUtils.hasText(threadId)) {
            builder.threadId(threadId);
        }
        toMetadata(context).forEach(builder::addMetadata);
        return builder.build();
    }

    public AiObservationContext fromMetadata(Map<String, ?> metadata, String threadId) {
        if (metadata == null || metadata.isEmpty()) {
            return AiObservationContext.builder().sessionId(threadId).build();
        }
        return AiObservationContext.builder()
                .traceId(stringValue(metadata.get("traceId")))
                .requestId(stringValue(metadata.get("requestId")))
                .sessionId(StringUtils.hasText(stringValue(metadata.get("sessionId"))) ? stringValue(metadata.get("sessionId")) : threadId)
                .flowId(stringValue(metadata.get("flowId")))
                .scene(stringValue(metadata.get("scene")))
                .componentType(resolveComponentType(stringValue(metadata.get("componentType"))))
                .componentName(stringValue(metadata.get("componentName")))
                .agentName(stringValue(metadata.get("agentName")))
                .modelProvider(stringValue(metadata.get("modelProvider")))
                .modelName(stringValue(metadata.get("modelName")))
                .toolName(stringValue(metadata.get("toolName")))
                .build();
    }

    public String serviceName() {
        return serviceName;
    }

    public boolean isEnabled() {
        return properties.isEnabled();
    }

    public boolean isHttpHeaderPropagationEnabled() {
        return properties.getHttp().isPropagateHeaders();
    }

    public static String newObservationId() {
        return UUID.randomUUID().toString().replace("-", "");
    }

    private void applyContext(AiObservationContext context) {
        if (context == null) {
            return;
        }
        putMdc(MDC_SERVICE_NAME, serviceName);
        putMdc(MDC_TRACE_ID, context.traceId());
        putMdc(MDC_REQUEST_ID, context.requestId());
        putMdc(MDC_SESSION_ID, context.sessionId());
        putMdc(MDC_FLOW_ID, context.flowId());
        putMdc(MDC_AI_SCENE, context.scene());
        putMdc(MDC_COMPONENT_TYPE, context.componentType() == null ? null : context.componentType().name());
        putMdc(MDC_COMPONENT_NAME, context.componentName());
        putMdc(MDC_AGENT_NAME, context.agentName());
        putMdc(MDC_MODEL_PROVIDER, context.modelProvider());
        putMdc(MDC_MODEL_NAME, context.modelName());
        putMdc(MDC_TOOL_NAME, context.toolName());
    }

    private void restore(Map<String, String> previous) {
        MDC.clear();
        if (previous != null && !previous.isEmpty()) {
            MDC.setContextMap(previous);
        }
    }

    private void putMdc(String key, String value) {
        if (!StringUtils.hasText(key)) {
            return;
        }
        if (StringUtils.hasText(value)) {
            MDC.put(key, value);
            return;
        }
        MDC.remove(key);
    }

    private void put(Map<String, Object> map, String key, Object value) {
        if (value != null) {
            map.put(key, value);
        }
    }

    private String stringValue(Object value) {
        return value == null ? null : value.toString();
    }

    private AiComponentType resolveComponentType(String value) {
        if (!StringUtils.hasText(value)) {
            return null;
        }
        try {
            return AiComponentType.valueOf(value);
        } catch (IllegalArgumentException exception) {
            return null;
        }
    }
}
