package com.superbiz.agent.service;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.superbiz.agent.domain.entity.AgentStep;
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.DiagnosisSessionRepository;
import com.superbiz.agent.repository.ToolInvocationRepository;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

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

@Service
@RequiredArgsConstructor
public class DiagnosisTraceService {

    private static final TypeReference<Map<String, Object>> JSON_MAP_TYPE = new TypeReference<>() {
    };

    private final DiagnosisSessionRepository diagnosisSessionRepository;
    private final AgentStepRepository agentStepRepository;
    private final ToolInvocationRepository toolInvocationRepository;
    private final ObjectMapper objectMapper;

    public DiagnosisTraceResponse getTrace(String sessionId) {
        DiagnosisSession session = diagnosisSessionRepository.findBySessionId(sessionId)
                .orElseThrow(() -> new SessionNotFoundException(sessionId));
        List<AgentStep> steps = agentStepRepository.findBySessionIdOrderByStepIndex(sessionId);
        List<ToolInvocation> toolInvocations = toolInvocationRepository.findBySessionIdOrderByIdAsc(sessionId);

        return DiagnosisTraceResponse.builder()
                .session(toSessionTrace(session))
                .steps(steps.stream().map(this::toAgentStepTrace).toList())
                .toolInvocations(toolInvocations.stream().map(this::toToolInvocationTrace).toList())
                .summary(toSummary(session, steps, toolInvocations))
                .build();
    }

    private DiagnosisTraceResponse.SessionTrace toSessionTrace(DiagnosisSession session) {
        return DiagnosisTraceResponse.SessionTrace.builder()
                .id(session.getId())
                .sessionId(session.getSessionId())
                .query(session.getQuery())
                .status(session.getStatus())
                .agentFlow(session.getAgentFlow())
                .totalDurationMs(session.getTotalDurationMs())
                .totalTokenCount(session.getTotalTokenCount())
                .stepCount(session.getStepCount())
                .toolCallCount(session.getToolCallCount())
                .answer(session.getAnswer())
                .selfEvaluationRaw(session.getSelfEvaluation())
                .selfEvaluation(parseJsonObject(session.getSelfEvaluation()))
                .feedback(session.getFeedback())
                .createdAt(session.getCreatedAt())
                .updatedAt(session.getUpdatedAt())
                .build();
    }

    private DiagnosisTraceResponse.AgentStepTrace toAgentStepTrace(AgentStep step) {
        return DiagnosisTraceResponse.AgentStepTrace.builder()
                .id(step.getId())
                .sessionId(step.getSessionId())
                .stepIndex(step.getStepIndex())
                .agentName(step.getAgentName())
                .modelInput(step.getModelInput())
                .modelOutput(step.getModelOutput())
                .thought(step.getThought())
                .hasToolCall(step.getHasToolCall())
                .durationMs(step.getDurationMs())
                .tokenCount(step.getTokenCount())
                .createdAt(step.getCreatedAt())
                .build();
    }

    private DiagnosisTraceResponse.ToolInvocationTrace toToolInvocationTrace(ToolInvocation invocation) {
        return DiagnosisTraceResponse.ToolInvocationTrace.builder()
                .id(invocation.getId())
                .sessionId(invocation.getSessionId())
                .stepId(invocation.getStepId())
                .toolName(invocation.getToolName())
                .inputParamsRaw(invocation.getInputParams())
                .inputParams(parseJsonObject(invocation.getInputParams()))
                .outputPreview(invocation.getOutputPreview())
                .outputLength(invocation.getOutputLength())
                .retrievalLayer(invocation.getRetrievalLayer())
                .l0MatchCount(invocation.getL0MatchCount())
                .l1MatchCount(invocation.getL1MatchCount())
                .truncated(invocation.getIsTruncated())
                .relevanceLevel(invocation.getRelevanceLevel())
                .dedupReason(invocation.getDedupReason())
                .retrievalDetailsRaw(invocation.getRetrievalDetails())
                .retrievalDetails(parseJsonObject(invocation.getRetrievalDetails()))
                .durationMs(invocation.getDurationMs())
                .success(invocation.getSuccess())
                .errorMessage(invocation.getErrorMessage())
                .createdAt(invocation.getCreatedAt())
                .build();
    }

    private DiagnosisTraceResponse.TraceSummary toSummary(
            DiagnosisSession session,
            List<AgentStep> steps,
            List<ToolInvocation> toolInvocations
    ) {
        Map<String, Object> selfEvaluation = parseJsonObject(session.getSelfEvaluation());
        return DiagnosisTraceResponse.TraceSummary.builder()
                .persistedStepCount(defaultInt(session.getStepCount()))
                .returnedStepCount(steps.size())
                .persistedToolCallCount(defaultInt(session.getToolCallCount()))
                .returnedToolCallCount(toolInvocations.size())
                .hasVerifierEvaluation(selfEvaluation != null && selfEvaluation.containsKey("verifier_evaluation"))
                .hasFeedback(session.getFeedback() != null && !session.getFeedback().isBlank())
                .build();
    }

    private Map<String, Object> parseJsonObject(String json) {
        if (json == null || json.isBlank()) {
            return null;
        }
        try {
            return objectMapper.readValue(json, JSON_MAP_TYPE);
        } catch (Exception ignored) {
            return null;
        }
    }

    private int defaultInt(Integer value) {
        return value == null ? 0 : value;
    }
}