package com.vincent.rsf.server.ai.service.impl.chat;

import com.vincent.rsf.framework.exception.CoolException;
import com.vincent.rsf.server.ai.dto.AiChatToolEventDto;
import com.vincent.rsf.server.ai.service.AiCallLogService;
import com.vincent.rsf.server.ai.service.MountedToolCallback;
import com.vincent.rsf.server.ai.store.AiCachedToolResult;
import com.vincent.rsf.server.ai.store.AiToolResultStore;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;

@Component
@RequiredArgsConstructor
public class AiToolObservationService {

    private final AiSseEventPublisher aiSseEventPublisher;
    private final AiToolResultStore aiToolResultStore;
    private final AiCallLogService aiCallLogService;

    public ToolCallback[] wrapToolCallbacks(ToolCallback[] toolCallbacks, SseEmitter emitter, String requestId,
                                            Long sessionId, AtomicLong toolCallSequence,
                                            AtomicLong toolSuccessCount, AtomicLong toolFailureCount,
                                            Long callLogId, Long userId, Long tenantId,
                                            AiThinkingTraceEmitter thinkingTraceEmitter) {
        if (toolCallbacks == null || toolCallbacks.length == 0) {
            return toolCallbacks;
        }
        List<ToolCallback> wrappedCallbacks = new ArrayList<>();
        for (ToolCallback callback : toolCallbacks) {
            if (callback == null) {
                continue;
            }
            wrappedCallbacks.add(new ObservableToolCallback(callback, emitter, requestId, sessionId, toolCallSequence,
                    toolSuccessCount, toolFailureCount, callLogId, userId, tenantId, thinkingTraceEmitter));
        }
        return wrappedCallbacks.toArray(new ToolCallback[0]);
    }

    private String summarizeToolPayload(String content, int maxLength) {
        if (!StringUtils.hasText(content)) {
            return null;
        }
        String normalized = content.trim()
                .replace("\r", " ")
                .replace("\n", " ")
                .replaceAll("\\s+", " ");
        return normalized.length() > maxLength ? normalized.substring(0, maxLength) : normalized;
    }

    private class ObservableToolCallback implements ToolCallback {

        private final ToolCallback delegate;
        private final SseEmitter emitter;
        private final String requestId;
        private final Long sessionId;
        private final AtomicLong toolCallSequence;
        private final AtomicLong toolSuccessCount;
        private final AtomicLong toolFailureCount;
        private final Long callLogId;
        private final Long userId;
        private final Long tenantId;
        private final AiThinkingTraceEmitter thinkingTraceEmitter;

        private ObservableToolCallback(ToolCallback delegate, SseEmitter emitter, String requestId,
                                       Long sessionId, AtomicLong toolCallSequence,
                                       AtomicLong toolSuccessCount, AtomicLong toolFailureCount,
                                       Long callLogId, Long userId, Long tenantId,
                                       AiThinkingTraceEmitter thinkingTraceEmitter) {
            this.delegate = delegate;
            this.emitter = emitter;
            this.requestId = requestId;
            this.sessionId = sessionId;
            this.toolCallSequence = toolCallSequence;
            this.toolSuccessCount = toolSuccessCount;
            this.toolFailureCount = toolFailureCount;
            this.callLogId = callLogId;
            this.userId = userId;
            this.tenantId = tenantId;
            this.thinkingTraceEmitter = thinkingTraceEmitter;
        }

        @Override
        public org.springframework.ai.tool.definition.ToolDefinition getToolDefinition() {
            return delegate.getToolDefinition();
        }

        @Override
        public org.springframework.ai.tool.metadata.ToolMetadata getToolMetadata() {
            return delegate.getToolMetadata();
        }

        @Override
        public String call(String toolInput) {
            return call(toolInput, null);
        }

        @Override
        public String call(String toolInput, ToolContext toolContext) {
            String toolName = delegate.getToolDefinition() == null ? "unknown" : delegate.getToolDefinition().name();
            String mountName = delegate instanceof MountedToolCallback ? ((MountedToolCallback) delegate).getMountName() : null;
            String toolCallId = requestId + "-tool-" + toolCallSequence.incrementAndGet();
            long startedAt = System.currentTimeMillis();
            AiCachedToolResult cachedToolResult = aiToolResultStore.getToolResult(tenantId, requestId, toolName, toolInput);
            if (cachedToolResult != null) {
                aiSseEventPublisher.emitSafely(emitter, "tool_result", AiChatToolEventDto.builder()
                        .requestId(requestId)
                        .sessionId(sessionId)
                        .toolCallId(toolCallId)
                        .toolName(toolName)
                        .mountName(mountName)
                        .status(cachedToolResult.isSuccess() ? "COMPLETED" : "FAILED")
                        .inputSummary(summarizeToolPayload(toolInput, 400))
                        .outputSummary(summarizeToolPayload(cachedToolResult.getOutput(), 600))
                        .errorMessage(cachedToolResult.getErrorMessage())
                        .durationMs(0L)
                        .timestamp(System.currentTimeMillis())
                        .build());
                if (thinkingTraceEmitter != null) {
                    thinkingTraceEmitter.onToolResult(toolName, toolCallId, !cachedToolResult.isSuccess());
                }
                if (cachedToolResult.isSuccess()) {
                    toolSuccessCount.incrementAndGet();
                    aiCallLogService.saveMcpCallLog(callLogId, requestId, sessionId, toolCallId, mountName, toolName,
                            "COMPLETED", summarizeToolPayload(toolInput, 400), summarizeToolPayload(cachedToolResult.getOutput(), 600),
                            null, 0L, userId, tenantId);
                    return cachedToolResult.getOutput();
                }
                toolFailureCount.incrementAndGet();
                aiCallLogService.saveMcpCallLog(callLogId, requestId, sessionId, toolCallId, mountName, toolName,
                        "FAILED", summarizeToolPayload(toolInput, 400), null, cachedToolResult.getErrorMessage(),
                        0L, userId, tenantId);
                throw new CoolException(cachedToolResult.getErrorMessage());
            }
            if (thinkingTraceEmitter != null) {
                thinkingTraceEmitter.onToolStart(toolName, toolCallId);
            }
            aiSseEventPublisher.emitSafely(emitter, "tool_start", AiChatToolEventDto.builder()
                    .requestId(requestId)
                    .sessionId(sessionId)
                    .toolCallId(toolCallId)
                    .toolName(toolName)
                    .mountName(mountName)
                    .status("STARTED")
                    .inputSummary(summarizeToolPayload(toolInput, 400))
                    .timestamp(startedAt)
                    .build());
            try {
                String output = toolContext == null ? delegate.call(toolInput) : delegate.call(toolInput, toolContext);
                long durationMs = System.currentTimeMillis() - startedAt;
                aiSseEventPublisher.emitSafely(emitter, "tool_result", AiChatToolEventDto.builder()
                        .requestId(requestId)
                        .sessionId(sessionId)
                        .toolCallId(toolCallId)
                        .toolName(toolName)
                        .mountName(mountName)
                        .status("COMPLETED")
                        .inputSummary(summarizeToolPayload(toolInput, 400))
                        .outputSummary(summarizeToolPayload(output, 600))
                        .durationMs(durationMs)
                        .timestamp(System.currentTimeMillis())
                        .build());
                if (thinkingTraceEmitter != null) {
                    thinkingTraceEmitter.onToolResult(toolName, toolCallId, false);
                }
                aiToolResultStore.cacheToolResult(tenantId, requestId, toolName, toolInput, true, output, null);
                toolSuccessCount.incrementAndGet();
                aiCallLogService.saveMcpCallLog(callLogId, requestId, sessionId, toolCallId, mountName, toolName,
                        "COMPLETED", summarizeToolPayload(toolInput, 400), summarizeToolPayload(output, 600),
                        null, durationMs, userId, tenantId);
                return output;
            } catch (RuntimeException e) {
                long durationMs = System.currentTimeMillis() - startedAt;
                aiSseEventPublisher.emitSafely(emitter, "tool_error", AiChatToolEventDto.builder()
                        .requestId(requestId)
                        .sessionId(sessionId)
                        .toolCallId(toolCallId)
                        .toolName(toolName)
                        .mountName(mountName)
                        .status("FAILED")
                        .inputSummary(summarizeToolPayload(toolInput, 400))
                        .errorMessage(e.getMessage())
                        .durationMs(durationMs)
                        .timestamp(System.currentTimeMillis())
                        .build());
                if (thinkingTraceEmitter != null) {
                    thinkingTraceEmitter.onToolResult(toolName, toolCallId, true);
                }
                aiToolResultStore.cacheToolResult(tenantId, requestId, toolName, toolInput, false, null, e.getMessage());
                toolFailureCount.incrementAndGet();
                aiCallLogService.saveMcpCallLog(callLogId, requestId, sessionId, toolCallId, mountName, toolName,
                        "FAILED", summarizeToolPayload(toolInput, 400), null, e.getMessage(),
                        durationMs, userId, tenantId);
                throw e;
            }
        }
    }
}