package org.dromara.ai.provider;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import okhttp3.sse.EventSources;
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.io.IOException;
import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * OpenAI 提供商实现
 */
@Slf4j
public class OpenAiProvider extends AbstractAiProvider {

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

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

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

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

    @Override
    protected void doInitialize() {
        String apiKey = (String) configuration.getOrDefault("api-key", "");
        if (apiKey == null || apiKey.isBlank()) {
            throw new IllegalStateException("OpenAI api-key 未配置");
        }
        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 baseUrl = (String) configuration.getOrDefault("base-url", "https://api.openai.com");
        String apiKey = (String) configuration.get("api-key");
        String model = Optional.ofNullable(request.getModel())
                .filter(s -> !s.isBlank())
                .orElse((String) configuration.getOrDefault("default-model", "gpt-3.5-turbo"));

        Map<String, Object> payload = new HashMap<>();
        payload.put("model", model);
        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());

        Request httpRequest = new Request.Builder()
                .url(baseUrl + "/v1/chat/completions")
                .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 IOException("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) {
        String baseUrl = (String) configuration.getOrDefault("base-url", "https://api.openai.com");
        String apiKey = (String) configuration.get("api-key");
        String model = Optional.ofNullable(request.getModel())
                .filter(s -> !s.isBlank())
                .orElse((String) configuration.getOrDefault("default-model", "gpt-3.5-turbo"));

        Map<String, Object> payload = new HashMap<>();
        payload.put("model", model);
        payload.put("messages", toOpenAiMessages(request.getMessages()));
        payload.put("stream", true);
        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());

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

        StringBuilder contentBuilder = new StringBuilder();
        EventSourceListener listener = new EventSourceListener() {
            @Override
            public void onEvent(EventSource eventSource, String id, String type, String data) {
                if (data == null) return;
                if ("[DONE]".equals(data.trim())) {
                    AiChatResponse finalResp = new AiChatResponse()
                            .setSuccess(true)
                            .setModel(model)
                            .setChoices(List.of(new AiChatResponse.Choice()
                                    .setIndex(0)
                                    .setMessage(new AiChatRequest.AiMessage()
                                            .setRole("assistant")
                                            .setContent(contentBuilder.toString()))));
                    callback.onComplete(finalResp);
                    eventSource.cancel();
                    return;
                }
                try {
                    JsonNode node = objectMapper.readTree(data);
                    JsonNode choices = node.path("choices");
                    if (choices.isArray() && choices.size() > 0) {
                        JsonNode delta = choices.get(0).path("delta");
                        String chunk = delta.path("content").asText(null);
                        if (chunk != null) {
                            contentBuilder.append(chunk);
                            callback.onChunk(chunk);
                        }
                    }
                } catch (Exception e) {
                    callback.onError(e);
                }
            }

            @Override
            public void onFailure(EventSource eventSource, Throwable t, Response response) {
                if (response != null) {
                    log.error("OpenAI 流失败: {} {}", response.code(), response.message(), t);
                } else {
                    log.error("OpenAI 流失败", t);
                }
                callback.onError(t);
            }
        };

        EventSources.createFactory(httpClient).newEventSource(httpRequest, listener);
    }

    @Override
    public List<AiModelInfo> getSupportedModels() {
        List<String> supported = (List<String>) configuration.getOrDefault("supported-models", List.of("gpt-3.5-turbo", "gpt-4"));
        String provider = getProviderName();
        List<AiModelInfo> infos = new ArrayList<>();
        for (String m : supported) {
            infos.add(new AiModelInfo()
                    .setId(m)
                    .setName(m)
                    .setDisplayName(m)
                    .setDescription("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);
        resp.setId(root.path("id").asText(null));
        resp.setCreated(root.path("created").isNumber() ? root.path("created").asLong() : null);
        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;
    }
}

