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

import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.agent.hook.AgentHook;
import com.alibaba.cloud.ai.graph.agent.hook.HookPosition;
import com.infoepoch.pms.agent.observability.ai.AiComponentType;
import com.infoepoch.pms.agent.observability.ai.AiObservationContext;
import com.infoepoch.pms.agent.observability.ai.AiObservationLogger;
import com.infoepoch.pms.agent.observability.ai.AiObservationSupport;
import com.infoepoch.pms.agent.observability.ai.AiReadableLogSupport;
import org.springframework.util.StringUtils;

import java.util.Collections;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

/**
 * 通用 Agent 生命周期观测 hook。
 */
public class AiAgentLifecycleObservationHook extends AgentHook {

    private static final String HOOK_NAME = "aiAgentLifecycleObservationHook";

    private final AiObservationLogger logger;
    private final AiObservationSupport support;

    public AiAgentLifecycleObservationHook(AiObservationLogger logger, AiObservationSupport support) {
        this.logger = logger;
        this.support = support;
    }

    @Override
    public String getName() {
        return HOOK_NAME;
    }

    @Override
    public HookPosition[] getHookPositions() {
        return new HookPosition[]{HookPosition.BEFORE_AGENT, HookPosition.AFTER_AGENT};
    }

    @Override
    public CompletableFuture<Map<String, Object>> beforeAgent(OverAllState state, RunnableConfig runnableConfig) {
        AiObservationContext context = buildContext(runnableConfig);
        logger.info(
                AiAgentLifecycleObservationHook.class,
                AiReadableLogSupport.format(
                        context,
                        "Agent开始",
                        "场景=%s，agent=%s".formatted(
                                safeValue(context.scene()),
                                safeValue(context.agentName())
                        )
                ),
                context,
                Map.of("status", "started")
        );
        return CompletableFuture.completedFuture(Collections.emptyMap());
    }

    @Override
    public CompletableFuture<Map<String, Object>> afterAgent(OverAllState state, RunnableConfig runnableConfig) {
        AiObservationContext context = buildContext(runnableConfig);
        logger.info(
                AiAgentLifecycleObservationHook.class,
                AiReadableLogSupport.format(
                        context,
                        "Agent完成",
                        "场景=%s，agent=%s".formatted(
                                safeValue(context.scene()),
                                safeValue(context.agentName())
                        )
                ),
                context,
                Map.of("status", "finished")
        );
        return CompletableFuture.completedFuture(Collections.emptyMap());
    }

    private AiObservationContext buildContext(RunnableConfig runnableConfig) {
        Map<String, Object> metadata = runnableConfig.metadata().orElse(Map.of());
        AiObservationContext current = support.fromMetadata(metadata, runnableConfig.threadId().orElse(null));
        return current.toBuilder()
                .componentType(AiComponentType.AGENT)
                .componentName(StringUtils.hasText(current.componentName()) ? current.componentName() : getAgentName())
                .agentName(StringUtils.hasText(current.agentName()) ? current.agentName() : getAgentName())
                .build();
    }

    private String safeValue(String value) {
        return StringUtils.hasText(value) ? value : "-";
    }
}
