package com.superbiz.agent.service;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.superbiz.agent.domain.entity.AgentStep;
import com.superbiz.agent.domain.entity.ChatSession;
import com.superbiz.agent.domain.entity.DiagnosisRun;
import com.superbiz.agent.domain.entity.DiagnosisSession;
import com.superbiz.agent.domain.entity.ToolInvocation;
import com.superbiz.agent.dto.DiagnosisTraceResponse;
import com.superbiz.agent.exception.SessionNotFoundException;
import com.superbiz.agent.repository.AgentStepRepository;
import com.superbiz.agent.repository.ChatSessionRepository;
import com.superbiz.agent.repository.DiagnosisRunRepository;
import com.superbiz.agent.repository.DiagnosisSessionRepository;
import com.superbiz.agent.repository.ToolInvocationRepository;
import org.junit.jupiter.api.Test;

import java.time.LocalDateTime;
import java.util.List;
import java.util.Optional;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
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.verifyNoInteractions;
import static org.mockito.Mockito.when;

class DiagnosisTraceServiceTest {

    private final DiagnosisSessionRepository diagnosisSessionRepository = mock(DiagnosisSessionRepository.class);
    private final ChatSessionRepository chatSessionRepository = mock(ChatSessionRepository.class);
    private final DiagnosisRunRepository diagnosisRunRepository = mock(DiagnosisRunRepository.class);
    private final AgentStepRepository agentStepRepository = mock(AgentStepRepository.class);
    private final ToolInvocationRepository toolInvocationRepository = mock(ToolInvocationRepository.class);
    private final DiagnosisTraceService service = new DiagnosisTraceService(
            diagnosisSessionRepository,
            chatSessionRepository,
            diagnosisRunRepository,
            agentStepRepository,
            toolInvocationRepository,
            new ObjectMapper()
    );

    @Test
    void getTraceWithoutRunIdResolvesLatestRun() {
        String sessionId = "trace-session-latest";
        LocalDateTime base = LocalDateTime.of(2026, 7, 10, 10, 0);
        DiagnosisRun latest = run(2L, sessionId, "run-latest", "second question", base.plusMinutes(1));
        AgentStep step = step(20L, sessionId, "run-latest", 0, "composer", base.plusMinutes(1));
        ToolInvocation invocation = invocation(30L, sessionId, "run-latest", "query_metrics", base.plusMinutes(1));

        when(diagnosisRunRepository.findFirstBySessionIdOrderByCreatedAtDescIdDesc(sessionId))
                .thenReturn(Optional.of(latest));
        when(chatSessionRepository.findBySessionId(sessionId)).thenReturn(Optional.of(chatSession(sessionId)));
        when(agentStepRepository.findByRunIdOrderByStepIndex("run-latest")).thenReturn(List.of(step));
        when(toolInvocationRepository.findByRunIdOrderByIdAsc("run-latest")).thenReturn(List.of(invocation));

        DiagnosisTraceResponse response = service.getTrace(sessionId);

        assertEquals("run-latest", response.getRunId());
        assertEquals("run-latest", response.getRun().getRunId());
        assertEquals("run-latest", response.getSummary().getResolvedRunId());
        assertEquals("second question", response.getSession().getQuery());
        assertEquals("run-latest", response.getSteps().get(0).getRunId());
        assertEquals("run-latest", response.getToolInvocations().get(0).getRunId());
        assertEquals(1, response.getChatSession().getMessagePairCount());
    }

    @Test
    void getTraceWithRunIdReturnsExactFirstRun() {
        String sessionId = "trace-session-exact";
        DiagnosisRun first = run(1L, sessionId, "run-first", "first question",
                LocalDateTime.of(2026, 7, 10, 10, 0));
        AgentStep firstStep = step(10L, sessionId, "run-first", 0, "planner", first.getCreatedAt());

        when(diagnosisRunRepository.findBySessionIdAndRunId(sessionId, "run-first"))
                .thenReturn(Optional.of(first));
        when(agentStepRepository.findByRunIdOrderByStepIndex("run-first")).thenReturn(List.of(firstStep));
        when(toolInvocationRepository.findByRunIdOrderByIdAsc("run-first")).thenReturn(List.of());

        DiagnosisTraceResponse response = service.getTrace(sessionId, "run-first");

        assertEquals("run-first", response.getRunId());
        assertEquals("first question", response.getRun().getQuery());
        assertEquals(List.of("planner"),
                response.getSteps().stream().map(DiagnosisTraceResponse.AgentStepTrace::getAgentName).toList());
        verify(diagnosisRunRepository, never()).findFirstBySessionIdOrderByCreatedAtDescIdDesc(any());
    }

    @Test
    void getTraceWithRunIdReturnsExactSecondRun() {
        String sessionId = "trace-session-exact";
        DiagnosisRun second = run(2L, sessionId, "run-second", "second question",
                LocalDateTime.of(2026, 7, 10, 10, 1));
        ToolInvocation secondTool = invocation(20L, sessionId, "run-second", "query_logs", second.getCreatedAt());

        when(diagnosisRunRepository.findBySessionIdAndRunId(sessionId, "run-second"))
                .thenReturn(Optional.of(second));
        when(agentStepRepository.findByRunIdOrderByStepIndex("run-second")).thenReturn(List.of());
        when(toolInvocationRepository.findByRunIdOrderByIdAsc("run-second")).thenReturn(List.of(secondTool));

        DiagnosisTraceResponse response = service.getTrace(sessionId, "run-second");

        assertEquals("run-second", response.getRunId());
        assertEquals("second question", response.getSession().getQuery());
        assertEquals(List.of("query_logs"),
                response.getToolInvocations().stream().map(DiagnosisTraceResponse.ToolInvocationTrace::getToolName).toList());
    }

    @Test
    void getTraceRejectsRunFromAnotherSession() {
        String runId = "run-other-session";
        when(diagnosisRunRepository.findBySessionIdAndRunId("path-session", runId)).thenReturn(Optional.empty());
        when(diagnosisRunRepository.findByRunId(runId)).thenReturn(Optional.of(run(
                1L, "actual-session", runId, "query", LocalDateTime.of(2026, 7, 10, 10, 0))));

        IllegalArgumentException error = assertThrows(IllegalArgumentException.class,
                () -> service.getTrace("path-session", runId));

        assertTrue(error.getMessage().contains("does not belong"));
        verifyNoInteractions(agentStepRepository, toolInvocationRepository);
    }

    @Test
    void getTraceThrowsWhenSessionMissing() {
        String sessionId = "missing-session";
        when(diagnosisRunRepository.findFirstBySessionIdOrderByCreatedAtDescIdDesc(sessionId))
                .thenReturn(Optional.empty());
        when(diagnosisSessionRepository.findBySessionId(sessionId)).thenReturn(Optional.empty());

        assertThrows(SessionNotFoundException.class, () -> service.getTrace(sessionId));

        verify(diagnosisRunRepository).findFirstBySessionIdOrderByCreatedAtDescIdDesc(sessionId);
        verify(diagnosisSessionRepository).findBySessionId(sessionId);
        verifyNoInteractions(agentStepRepository, toolInvocationRepository);
    }

    @Test
    void getTraceIsReadOnly() {
        String sessionId = "trace-session-readonly";
        DiagnosisRun latest = run(3L, sessionId, "run-readonly", "readonly question",
                LocalDateTime.of(2026, 7, 10, 10, 2));

        when(diagnosisRunRepository.findFirstBySessionIdOrderByCreatedAtDescIdDesc(sessionId))
                .thenReturn(Optional.of(latest));
        when(agentStepRepository.findByRunIdOrderByStepIndex("run-readonly")).thenReturn(List.of());
        when(toolInvocationRepository.findByRunIdOrderByIdAsc("run-readonly")).thenReturn(List.of());

        service.getTrace(sessionId);

        verify(chatSessionRepository, never()).save(any());
        verify(diagnosisRunRepository, never()).save(any());
        verify(diagnosisSessionRepository, never()).save(any());
        verify(agentStepRepository, never()).save(any());
        verify(toolInvocationRepository, never()).save(any());
    }

    @Test
    void listRunSummariesDoesNotExpandTraceDetails() {
        String sessionId = "trace-session-runs";
        DiagnosisRun second = run(2L, sessionId, "run-second", "second question",
                LocalDateTime.of(2026, 7, 10, 10, 1));
        second.setAnswer("answer ".repeat(40));
        DiagnosisRun first = run(1L, sessionId, "run-first", "first question",
                LocalDateTime.of(2026, 7, 10, 10, 0));

        when(diagnosisRunRepository.findBySessionIdOrderByCreatedAtDescIdDesc(sessionId))
                .thenReturn(List.of(second, first));

        List<DiagnosisTraceResponse.RunSummary> summaries = service.listRunSummaries(sessionId);

        assertEquals(List.of("run-second", "run-first"),
                summaries.stream().map(DiagnosisTraceResponse.RunSummary::getRunId).toList());
        assertEquals("second question", summaries.get(0).getQuery());
        assertTrue(summaries.get(0).getAnswerPreview().length() <= 160);
        verifyNoInteractions(agentStepRepository, toolInvocationRepository);
    }

    @Test
    void legacyTraceFallbackKeepsHistoricalSessionReadable() {
        String sessionId = "legacy-session";
        DiagnosisSession legacy = DiagnosisSession.builder()
                .id(1L)
                .sessionId(sessionId)
                .query("legacy question")
                .status("SUCCESS")
                .stepCount(0)
                .toolCallCount(0)
                .createdAt(LocalDateTime.of(2026, 7, 10, 9, 0))
                .updatedAt(LocalDateTime.of(2026, 7, 10, 9, 1))
                .build();

        when(diagnosisRunRepository.findFirstBySessionIdOrderByCreatedAtDescIdDesc(sessionId))
                .thenReturn(Optional.empty());
        when(diagnosisSessionRepository.findBySessionId(sessionId)).thenReturn(Optional.of(legacy));
        when(agentStepRepository.findBySessionId(sessionId)).thenReturn(List.of());
        when(toolInvocationRepository.findBySessionIdOrderByIdAsc(sessionId)).thenReturn(List.of());

        DiagnosisTraceResponse response = service.getTrace(sessionId);

        assertNull(response.getRunId());
        assertNull(response.getRun());
        assertEquals("legacy question", response.getSession().getQuery());
    }

    private DiagnosisRun run(Long id, String sessionId, String runId, String query, LocalDateTime createdAt) {
        return DiagnosisRun.builder()
                .id(id)
                .sessionId(sessionId)
                .runId(runId)
                .query(query)
                .status("SUCCESS")
                .agentFlow("CHAT")
                .answer("answer for " + runId)
                .selfEvaluation("{\"verifier_evaluation\":{\"verdict\":\"PASS\"}}")
                .feedback("useful")
                .stepCount(1)
                .toolCallCount(1)
                .createdAt(createdAt)
                .updatedAt(createdAt.plusSeconds(1))
                .build();
    }

    private ChatSession chatSession(String sessionId) {
        return ChatSession.builder()
                .id(99L)
                .sessionId(sessionId)
                .status("ACTIVE")
                .messagePairCount(1)
                .createdAt(LocalDateTime.of(2026, 7, 10, 9, 0))
                .lastActiveAt(LocalDateTime.of(2026, 7, 10, 10, 0))
                .build();
    }

    private AgentStep step(Long id, String sessionId, String runId, int stepIndex, String agentName, LocalDateTime createdAt) {
        return AgentStep.builder()
                .id(id)
                .sessionId(sessionId)
                .runId(runId)
                .stepIndex(stepIndex)
                .agentName(agentName)
                .createdAt(createdAt)
                .build();
    }

    private ToolInvocation invocation(Long id, String sessionId, String runId, String toolName, LocalDateTime createdAt) {
        return ToolInvocation.builder()
                .id(id)
                .sessionId(sessionId)
                .runId(runId)
                .toolName(toolName)
                .inputParams("{\"query\":\"timeout\"}")
                .retrievalDetails("{\"evidence_refs\":[]}")
                .success(true)
                .createdAt(createdAt)
                .build();
    }
}