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

import ch.qos.logback.classic.Logger;
import ch.qos.logback.classic.spi.ILoggingEvent;
import ch.qos.logback.core.read.ListAppender;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.infoepoch.pms.agent.observability.ai.hook.AiAgentLifecycleObservationHook;
import com.infoepoch.pms.agent.observability.ai.hook.AiModelMessageObservationHook;
import com.infoepoch.pms.agent.properties.AiObservabilityProperties;
import org.junit.jupiter.api.Test;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;

import java.util.List;
import java.util.stream.Collectors;

import static org.junit.jupiter.api.Assertions.assertTrue;

class AiObservationHookTest {

    @Test
    void shouldLogReadableAgentLifecycleMessages() {
        AiObservationSupport support = new AiObservationSupport(new AiObservabilityProperties(), "pms-agent");
        AiObservationLogger logger = new AiObservationLogger(support);
        AiAgentLifecycleObservationHook hook = new AiAgentLifecycleObservationHook(logger, support);
        ListAppender<ILoggingEvent> appender = attachAppender(AiAgentLifecycleObservationHook.class);

        hook.beforeAgent(null, runnableConfig(support));
        hook.afterAgent(null, runnableConfig(support));

        String logs = joinedLogs(appender);
        assertTrue(logs.contains("[链路ID=trace-1]"));
        assertTrue(logs.contains("[会话ID=session-1]"));
        assertTrue(logs.contains("[阶段=Agent开始]"));
        assertTrue(logs.contains("[阶段=Agent完成]"));
        assertTrue(logs.contains("场景=通用对话"));
        assertTrue(logs.contains("agent=careAgent"));
    }

    @Test
    void shouldLogReadableModelRequestAndResponseMessages() {
        AiObservationSupport support = new AiObservationSupport(new AiObservabilityProperties(), "pms-agent");
        AiObservationLogger logger = new AiObservationLogger(support);
        AiModelMessageObservationHook hook = new AiModelMessageObservationHook(logger, support);
        ListAppender<ILoggingEvent> appender = attachAppender(AiModelMessageObservationHook.class);
        List<Message> beforeMessages = List.of(
                new SystemMessage("你是精准关爱智能体"),
                new UserMessage("活动历史有哪些")
        );
        List<Message> afterMessages = List.of(
                new SystemMessage("你是精准关爱智能体"),
                new UserMessage("活动历史有哪些"),
                new AssistantMessage("这里是模型回复")
        );

        hook.beforeModel(beforeMessages, runnableConfig(support));
        hook.afterModel(afterMessages, runnableConfig(support));

        String logs = joinedLogs(appender);
        assertTrue(logs.contains("[链路ID=trace-1]"));
        assertTrue(logs.contains("[会话ID=session-1]"));
        assertTrue(logs.contains("[阶段=通用对话模型请求]"));
        assertTrue(logs.contains("[阶段=通用对话模型完整返回]"));
        assertTrue(logs.contains("role=system"));
        assertTrue(logs.contains("role=user"));
        assertTrue(logs.contains("活动历史有哪些"));
        assertTrue(logs.contains("这里是模型回复"));
    }

    private RunnableConfig runnableConfig(AiObservationSupport support) {
        AiObservationContext context = AiObservationContext.builder()
                .traceId("trace-1")
                .requestId("request-1")
                .sessionId("session-1")
                .flowId("flow-1")
                .scene("通用对话")
                .componentName("careAgent")
                .agentName("careAgent")
                .build();
        return support.buildRunnableConfig("session-1", context);
    }

    private ListAppender<ILoggingEvent> attachAppender(Class<?> loggerClass) {
        Logger logger = (Logger) LoggerFactory.getLogger(loggerClass);
        ListAppender<ILoggingEvent> appender = new ListAppender<>();
        appender.start();
        logger.addAppender(appender);
        return appender;
    }

    private String joinedLogs(ListAppender<ILoggingEvent> appender) {
        return appender.list.stream()
                .map(ILoggingEvent::getFormattedMessage)
                .collect(Collectors.joining("\n"));
    }
}
