package org.dromara.ai.provider.impl;

import org.dromara.ai.domain.dto.AiChatRequest;
import org.dromara.ai.domain.dto.AiChatResponse;
import org.dromara.ai.domain.dto.AiModelInfo;
import org.dromara.ai.domain.enums.AiProviderType;
import org.dromara.ai.provider.AbstractAiProvider;
import org.dromara.ai.provider.AiProvider;
import org.springframework.stereotype.Component;

import java.time.Instant;
import java.util.ArrayList;
import java.util.List;

@Component
public class MockAiProvider extends AbstractAiProvider {

    @Override
    public AiProviderType getProviderType() {
        return AiProviderType.CUSTOM;
    }

    @Override
    public String getProviderName() {
        return "mock";
    }

    @Override
    public String getVersion() {
        return "1.0.0";
    }

    @Override
    protected void doInitialize() {
    }

    @Override
    protected void doDestroy() {
    }

    @Override
    protected AiChatResponse doChat(AiChatRequest request) {
        String content = buildReplyContent(request);
        AiChatRequest.AiMessage message = new AiChatRequest.AiMessage()
                .setRole("assistant")
                .setContent(content);
        AiChatResponse.Choice choice = new AiChatResponse.Choice()
                .setIndex(0)
                .setMessage(message)
                .setFinishReason(AiChatResponse.FinishReason.STOP.getValue());
        AiChatResponse.Usage usage = new AiChatResponse.Usage()
                .setPromptTokens(content.length())
                .setCompletionTokens(content.length())
                .setTotalTokens(content.length() * 2);
        return new AiChatResponse()
                .setId("mock-" + Instant.now().toEpochMilli())
                .setObject("chat.completion")
                .setCreated(Instant.now().getEpochSecond())
                .setModel(request.getModel() != null ? request.getModel() : "mock-gpt")
                .setChoices(List.of(choice))
                .setUsage(usage)
                .setSuccess(true);
    }

    @Override
    protected void doChatStream(AiChatRequest request, AiProvider.StreamCallback callback) {
        String content = buildReplyContent(request);
        for (int i = 0; i < content.length(); i++) {
            String c = String.valueOf(content.charAt(i));
            String chunk = "{\"choices\":[{\"delta\":{\"content\":\"" + escapeJson(c) + "\"}}]}";
            callback.onChunk(chunk);
        }
        AiChatRequest.AiMessage message = new AiChatRequest.AiMessage()
                .setRole("assistant")
                .setContent(content);
        AiChatResponse.Choice choice = new AiChatResponse.Choice()
                .setIndex(0)
                .setMessage(message)
                .setFinishReason(AiChatResponse.FinishReason.STOP.getValue());
        AiChatResponse response = new AiChatResponse()
                .setId("mock-" + Instant.now().toEpochMilli())
                .setObject("chat.completion")
                .setCreated(Instant.now().getEpochSecond())
                .setModel(request.getModel() != null ? request.getModel() : "mock-gpt")
                .setChoices(List.of(choice))
                .setUsage(new AiChatResponse.Usage()
                        .setPromptTokens(content.length())
                        .setCompletionTokens(content.length())
                        .setTotalTokens(content.length() * 2))
                .setSuccess(true);
        callback.onComplete(response);
    }

    @Override
    public List<AiModelInfo> getSupportedModels() {
        List<AiModelInfo> list = new ArrayList<>();
        AiModelInfo.ModelParameters params1 = new AiModelInfo.ModelParameters()
                .setDefaultMaxTokens(8192);
        list.add(new AiModelInfo()
                .setId("mock-gpt")
                .setName("mock-gpt")
                .setProvider(getProviderName())
                .setDescription("Mock GPT")
                .setMaxContextLength(8192)
                .setSupportStream(true)
                .setStatus(AiModelInfo.ModelStatus.AVAILABLE)
                .setParameters(params1));
        AiModelInfo.ModelParameters params2 = new AiModelInfo.ModelParameters()
                .setDefaultMaxTokens(4096);
        list.add(new AiModelInfo()
                .setId("gpt-3.5-turbo")
                .setName("gpt-3.5-turbo")
                .setProvider(getProviderName())
                .setDescription("Mock GPT-3.5")
                .setMaxContextLength(4096)
                .setSupportStream(true)
                .setStatus(AiModelInfo.ModelStatus.AVAILABLE)
                .setParameters(params2));
        return list;
    }

    private String buildReplyContent(AiChatRequest request) {
        String last = "";
        if (request.getMessages() != null && !request.getMessages().isEmpty()) {
            for (AiChatRequest.AiMessage m : request.getMessages()) {
                if ("user".equals(m.getRole())) {
                    last = m.getContent();
                }
            }
        }
        String prefix = "这是AI模拟回复：";
        return prefix + (last != null ? last : "");
    }

    private String escapeJson(String s) {
        return s.replace("\\", "\\\\").replace("\"", "\\\"").replace("\n", "\\n").replace("\r", "\\r");
    }
}
