package org.dromara.ai.provider;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.*;
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 java.util.*;
import java.util.concurrent.TimeUnit;

public class FastGptProvider extends AbstractAiProvider {

    private OkHttpClient httpClient;
    private final ObjectMapper objectMapper = new ObjectMapper();

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

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

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

    @Override
    protected void doInitialize() {
        String apiKey = (String) configuration.getOrDefault("api-key", "");
        String baseUrl = (String) configuration.getOrDefault("base-url", "");
        if (apiKey == null || apiKey.isBlank() || baseUrl == null || baseUrl.isBlank()) {
            throw new IllegalStateException("FastGPT 配置不完整");
        }
        Integer timeoutMs = (Integer) configuration.getOrDefault("timeout-ms", 30000);
        this.httpClient = new OkHttpClient.Builder()
                .connectTimeout(timeoutMs, TimeUnit.MILLISECONDS)
                .readTimeout(timeoutMs, TimeUnit.MILLISECONDS)
                .writeTimeout(timeoutMs, TimeUnit.MILLISECONDS)
                .build();
    }

    @Override
    protected void doDestroy() {
        this.httpClient = null;
    }

    @Override
    protected AiChatResponse doChat(AiChatRequest request) {
        String apiKey = (String) configuration.get("api-key");
        String baseUrl = (String) configuration.get("base-url");
        String appId = (String) configuration.getOrDefault("default-app-id", "");

        Map<String, Object> payload = new HashMap<>();
        payload.put("appId", appId);
        payload.put("messages", toOpenAiMessages(request.getMessages()));

        Request httpRequest = new Request.Builder()
                .url(baseUrl + "/api/openapi/chat")
                .header("Authorization", "Bearer " + apiKey)
                .header("Content-Type", "application/json")
                .post(RequestBody.create(MediaType.parse("application/json"), writeJson(payload)))
                .build();

        try (Response resp = httpClient.newCall(httpRequest).execute()) {
            if (!resp.isSuccessful()) {
                throw new RuntimeException("FastGPT 响应失败: " + resp.code() + " " + resp.message());
            }
            String body = Objects.requireNonNull(resp.body()).string();
            JsonNode root = objectMapper.readTree(body);
            return toAiChatResponse(root);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    protected void doChatStream(AiChatRequest request, StreamCallback callback) {
        AiChatResponse resp = chat(request);
        String content = resp.getChoices() != null && !resp.getChoices().isEmpty()
                ? Optional.ofNullable(resp.getChoices().get(0).getMessage().getContent()).orElse("")
                : "";
        int step = Math.max(1, content.length() / 20);
        for (int i = 0; i < content.length(); i += step) {
            int end = Math.min(content.length(), i + step);
            String chunk = content.substring(i, end);
            if (!chunk.isEmpty()) {
                callback.onChunk(chunk);
            }
        }
        callback.onComplete(resp);
    }

    @Override
    public List<AiModelInfo> getSupportedModels() {
        List<String> supported = (List<String>) configuration.getOrDefault("supported-models", List.of("fastgpt"));
        String provider = getProviderName();
        List<AiModelInfo> infos = new ArrayList<>();
        for (String m : supported) {
            infos.add(new AiModelInfo()
                    .setId(m)
                    .setName(m)
                    .setDisplayName(m)
                    .setDescription("FastGPT model " + m)
                    .setProvider(provider)
                    .setSupportStream(true)
                    .setAvailable(true)
                    .setType(AiModelInfo.ModelType.CHAT)
            );
        }
        return infos;
    }

    private List<Map<String, Object>> toOpenAiMessages(List<AiChatRequest.AiMessage> messages) {
        List<Map<String, Object>> list = new ArrayList<>();
        if (messages == null) return list;
        for (AiChatRequest.AiMessage msg : messages) {
            Map<String, Object> item = new HashMap<>();
            item.put("role", msg.getRole());
            item.put("content", msg.getContent());
            list.add(item);
        }
        return list;
    }

    private String writeJson(Object o) {
        try {
            return objectMapper.writeValueAsString(o);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private AiChatResponse toAiChatResponse(JsonNode root) {
        AiChatResponse resp = new AiChatResponse();
        resp.setSuccess(true);
        List<AiChatResponse.Choice> choices = new ArrayList<>();
        String content = root.path("data").path("text").asText("");
        AiChatRequest.AiMessage message = new AiChatRequest.AiMessage()
                .setRole("assistant")
                .setContent(content);
        choices.add(new AiChatResponse.Choice().setIndex(0).setMessage(message));
        resp.setChoices(choices);
        resp.setModel(getProviderName());
        return resp;
    }
}
