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

import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.agent.hook.HookPosition;
import com.alibaba.cloud.ai.graph.agent.hook.messages.AgentCommand;
import com.alibaba.cloud.ai.graph.agent.hook.messages.MessagesModelHook;
import com.alibaba.cloud.ai.graph.agent.hook.messages.UpdatePolicy;
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.ai.chat.messages.Message;
import org.springframework.util.StringUtils;

import java.util.List;
import java.util.Map;

/**
 * 通用模型消息观测 hook。
 */
public class AiModelMessageObservationHook extends MessagesModelHook {

    private static final String HOOK_NAME = "aiModelMessageObservationHook";

    private final AiObservationLogger logger;
    private final AiObservationSupport support;

    public AiModelMessageObservationHook(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_MODEL, HookPosition.AFTER_MODEL};
    }

    @Override
    public AgentCommand beforeModel(List<Message> messages, RunnableConfig runnableConfig) {
        AiObservationContext context = buildContext(runnableConfig);
        String renderedMessages = AiReadableLogSupport.formatModelMessages(
                context,
                AiReadableLogSupport.stageWithScene(context, "模型请求"),
                messages
        );
        logger.info(
                AiModelMessageObservationHook.class,
                renderedMessages,
                context,
                Map.of(
                        "status", "started",
                        "ai.input.preview", renderedMessages,
                        "ai.input.length", renderedMessages.length()
                )
        );
        return new AgentCommand(messages, UpdatePolicy.REPLACE);
    }

    @Override
    public AgentCommand afterModel(List<Message> messages, RunnableConfig runnableConfig) {
        AiObservationContext context = buildContext(runnableConfig);
        String response = AiReadableLogSupport.findLastAssistantText(messages);
        logger.info(
                AiModelMessageObservationHook.class,
                AiReadableLogSupport.formatModelResponse(
                        context,
                        AiReadableLogSupport.stageWithScene(context, "模型完整返回"),
                        response
                ),
                context,
                Map.of(
                        "status", "finished",
                        "ai.output.preview", response == null ? "" : response,
                        "ai.output.length", response == null ? 0 : response.length()
                )
        );
        return new AgentCommand(messages, UpdatePolicy.REPLACE);
    }

    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.MODEL)
                .componentName(StringUtils.hasText(current.componentName())
                        ? current.componentName()
                        : getAgentName() + ".model")
                .agentName(StringUtils.hasText(current.agentName()) ? current.agentName() : getAgentName())
                .build();
    }
}
