package com.superbiz.agent.harness.audit;

import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.superbiz.agent.domain.entity.AgentStep;
import com.superbiz.agent.harness.agent.DiagnosisAgentFactory;
import com.superbiz.agent.repository.AgentStepRepository;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.UserMessage;

import java.util.List;
import java.util.Optional;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

class HarnessAgentAuditHookTest {

    @Test
    void persistsExactIdentityAndMetadataWithoutContentOrArguments() {
        AgentStepRepository repository = mock(AgentStepRepository.class);
        AgentStep persisted = AgentStep.builder().id(7L).build();
        when(repository.save(any(AgentStep.class))).thenReturn(persisted);
        when(repository.findById(7L)).thenReturn(Optional.of(persisted));
        HarnessAgentAuditHook hook = new HarnessAgentAuditHook(
                repository, new ObjectMapper(), DiagnosisAgentFactory.AGENT_NAME);
        RunnableConfig config = RunnableConfig.builder()
                .addMetadata("sessionId", "session-audit")
                .addMetadata("runId", "run-audit")
                .build();

        hook.beforeModel(List.of(new UserMessage("secret-query")), config);
        AssistantMessage response = AssistantMessage.builder()
                .content("secret-model-output")
                .toolCalls(List.of(new AssistantMessage.ToolCall(
                        "call-1", "function", "query_logs", "{\"query\":\"secret-argument\"}")))
                .build();
        hook.afterModel(List.of(response), config);

        ArgumentCaptor<AgentStep> captor = ArgumentCaptor.forClass(AgentStep.class);
        verify(repository, org.mockito.Mockito.times(2)).save(captor.capture());
        AgentStep started = captor.getAllValues().get(0);
        AgentStep completed = captor.getAllValues().get(1);
        assertEquals("session-audit", started.getSessionId());
        assertEquals("run-audit", started.getRunId());
        assertFalse(started.getModelInput().contains("secret-query"));
        assertFalse(completed.getModelOutput().contains("secret-model-output"));
        assertFalse(completed.getModelOutput().contains("secret-argument"));
        assertEquals("{\"has_text\":true,\"tool_names\":[\"query_logs\"]}", completed.getModelOutput());
        assertNull(completed.getThought());
    }

    @Test
    void missingIdentitySkipsPersistence() {
        AgentStepRepository repository = mock(AgentStepRepository.class);
        HarnessAgentAuditHook hook = new HarnessAgentAuditHook(
                repository, new ObjectMapper(), DiagnosisAgentFactory.AGENT_NAME);

        hook.beforeModel(List.of(new UserMessage("secret")), RunnableConfig.builder().build());

        verify(repository, never()).save(any());
    }
}