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 AzureOpenAiProvider extends AbstractAiProvider {

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

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

    @Override
    public String getProviderName() {
        return "azure-openai";
    }

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

    @Override
    protected void doInitialize() {
        String apiKey = (String) configuration.getOrDefault("api-key", "");
        String endpoint = (String) configuration.getOrDefault("endpoint", "");
        String deployment = (String) configuration.getOrDefault("deployment-name", "");
        if (apiKey == null || apiKey.isBlank() || endpoint == null || endpoint.isBlank() || deployment == null || deployment.isBlank()) {
            throw new IllegalStateException("Azure OpenAI 配置不完整");
        }
        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 endpoint = (String) configuration.get("endpoint");
        String deployment = (String) configuration.get("deployment-name");
        String apiVersion = (String) configuration.getOrDefault("api-version", "2023-12-01-preview");

        String model = Optional.ofNullable(request.getModel())
                .filter(s -> !s.isBlank())
                .orElse((String) configuration.getOrDefault("default-model", deployment));

        Map<String, Object> payload = new HashMap<>();
        payload.put("messages", toOpenAiMessages(request.getMessages()));
        if (request.getTemperature() != null) payload.put("temperature", request.getTemperature());
        if (request.getMaxTokens() != null) payload.put("max_tokens", request.getMaxTokens());
        if (request.getTopP() != null) payload.put("top_p", request.getTopP());
        if (request.getStop() != null && !request.getStop().isEmpty()) payload.put("stop", request.getStop());

        String url = endpoint + "/openai/deployments/" + deployment + "/chat/completions?api-version=" + apiVersion;
        Request httpRequest = new Request.Builder()
                .url(url)
                .header("api-key", 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("Azure OpenAI 响应失败: " + resp.code() + " " + resp.message());
            }
            String body = Objects.requireNonNull(resp.body()).string();
            JsonNode root = objectMapper.readTree(body);
            return toAiChatResponse(root, model);
        } 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());
        String provider = getProviderName();
        List<AiModelInfo> infos = new ArrayList<>();
        for (String m : supported) {
            infos.add(new AiModelInfo()
                    .setId(m)
                    .setName(m)
                    .setDisplayName(m)
                    .setDescription("Azure OpenAI 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, String model) {
        AiChatResponse resp = new AiChatResponse();
        resp.setSuccess(true);
        resp.setModel(model);
        List<AiChatResponse.Choice> choices = new ArrayList<>();
        JsonNode c = root.path("choices");
        if (c.isArray()) {
            for (int i = 0; i < c.size(); i++) {
                JsonNode ci = c.get(i);
                AiChatRequest.AiMessage message = new AiChatRequest.AiMessage()
                        .setRole(Optional.ofNullable(ci.path("message").path("role").asText(null)).orElse("assistant"))
                        .setContent(ci.path("message").path("content").asText(""));
                choices.add(new AiChatResponse.Choice()
                        .setIndex(ci.path("index").isNumber() ? ci.path("index").asInt() : i)
                        .setMessage(message)
                        .setFinishReason(ci.path("finish_reason").asText(null)));
            }
        }
        resp.setChoices(choices);
        JsonNode usage = root.path("usage");
        if (usage.isObject()) {
            resp.setUsage(new AiChatResponse.Usage()
                    .setPromptTokens(usage.path("prompt_tokens").isNumber() ? usage.path("prompt_tokens").asInt() : null)
                    .setCompletionTokens(usage.path("completion_tokens").isNumber() ? usage.path("completion_tokens").asInt() : null)
                    .setTotalTokens(usage.path("total_tokens").isNumber() ? usage.path("total_tokens").asInt() : null));
        }
        return resp;
    }
}
