package tech.easyflow.ai.invoke.provider.base;

import cn.hutool.core.util.StrUtil;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.MediaType;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;
import okhttp3.ResponseBody;
import okio.BufferedSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import tech.easyflow.ai.entity.Model;
import tech.easyflow.ai.invoke.exception.ModelInvokeException;
import tech.easyflow.ai.invoke.mapper.OpenAiProtocolMapper;
import tech.easyflow.ai.invoke.model.UnifiedChatChunk;
import tech.easyflow.ai.invoke.protocol.openai.OpenAiChatCompletionChunkResponse;
import tech.easyflow.common.util.OkHttpClientUtil;

import java.io.IOException;
import java.util.Objects;

public abstract class AbstractOpenAiCompatibleGateway {

    private static final Logger log = LoggerFactory.getLogger(AbstractOpenAiCompatibleGateway.class);
    private static final MediaType JSON_TYPE = MediaType.parse("application/json; charset=utf-8");

    protected final ObjectMapper objectMapper;
    protected final OpenAiProtocolMapper openAiProtocolMapper;
    protected final ThreadPoolTaskExecutor sseThreadPool;

    protected AbstractOpenAiCompatibleGateway(ObjectMapper objectMapper,
                                              OpenAiProtocolMapper openAiProtocolMapper,
                                              ThreadPoolTaskExecutor sseThreadPool) {
        this.objectMapper = objectMapper;
        this.openAiProtocolMapper = openAiProtocolMapper;
        this.sseThreadPool = sseThreadPool;
    }

    protected Response executePost(Model model, Object requestBody) {
        try {
            String body = objectMapper.writeValueAsString(requestBody);
            Request request = new Request.Builder()
                    .url(buildUrl(model))
                    .addHeader("Authorization", "Bearer " + model.checkAndGetApiKey())
                    .addHeader("Content-Type", "application/json")
                    .post(RequestBody.create(body, JSON_TYPE))
                    .build();
            return OkHttpClientUtil.buildDefaultClient().newCall(request).execute();
        } catch (ModelInvokeException e) {
            throw e;
        } catch (Exception e) {
            throw ModelInvokeException.badGateway("调用上游模型失败: " + e.getMessage());
        }
    }

    protected String buildUrl(Model model) {
        return StrUtil.removeSuffix(model.checkAndGetEndpoint(), "/")
                + "/"
                + StrUtil.removePrefix(model.checkAndGetRequestPath(), "/");
    }

    protected void validateResponse(Response response) throws IOException {
        if (response.isSuccessful()) {
            return;
        }
        String message = extractUpstreamErrorMessage(response);
        int code = response.code();
        if (code >= 500 || code == 429) {
            throw ModelInvokeException.serviceUnavailable(message);
        }
        throw ModelInvokeException.badGateway(message);
    }

    protected String extractUpstreamErrorMessage(Response response) throws IOException {
        ResponseBody body = response.body();
        String bodyString = body == null ? "" : body.string();
        if (StrUtil.isBlank(bodyString)) {
            return "上游模型调用失败，HTTP " + response.code();
        }
        try {
            JsonNode root = objectMapper.readTree(bodyString);
            String errorMessage = root.path("error").path("message").asText();
            if (StrUtil.isNotBlank(errorMessage)) {
                return "上游模型调用失败，HTTP " + response.code() + ": " + errorMessage;
            }
        } catch (Exception ignored) {
        }
        return "上游模型调用失败，HTTP " + response.code() + ": " + bodyString;
    }

    protected void streamResponse(Response response,
                                  tech.easyflow.ai.invoke.provider.UnifiedChatChunkObserver observer) {
        sseThreadPool.execute(() -> {
            try (Response closeableResponse = response; ResponseBody body = closeableResponse.body()) {
                if (body == null) {
                    throw new IOException("上游流式响应体为空");
                }
                BufferedSource source = body.source();
                while (!source.exhausted()) {
                    String line = source.readUtf8Line();
                    if (line == null) {
                        break;
                    }
                    if (line.isBlank() || !line.startsWith("data:")) {
                        continue;
                    }
                    String payload = line.substring(5).trim();
                    if (Objects.equals(payload, "[DONE]")) {
                        observer.onComplete();
                        return;
                    }
                    OpenAiChatCompletionChunkResponse chunkResponse = objectMapper.readValue(
                            payload,
                            OpenAiChatCompletionChunkResponse.class
                    );
                    UnifiedChatChunk chunk = openAiProtocolMapper.toUnifiedChunk(chunkResponse);
                    observer.onChunk(chunk);
                }
                observer.onComplete();
            } catch (Exception e) {
                log.error("streamResponse error: {}", e.getMessage(), e);
                observer.onError(e);
            }
        });
    }
}
