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

import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.Response;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;
import tech.easyflow.ai.entity.Model;
import tech.easyflow.ai.invoke.mapper.OpenAiProtocolMapper;
import tech.easyflow.ai.invoke.model.UnifiedChatRequest;
import tech.easyflow.ai.invoke.model.UnifiedChatResponse;
import tech.easyflow.ai.invoke.protocol.openai.OpenAiChatCompletionRequest;
import tech.easyflow.ai.invoke.protocol.openai.OpenAiChatCompletionResponse;
import tech.easyflow.ai.invoke.provider.ModelProviderGateway;
import tech.easyflow.ai.invoke.provider.UnifiedChatChunkObserver;
import tech.easyflow.ai.invoke.provider.base.AbstractOpenAiCompatibleGateway;

import java.util.Set;

@Component
public class OpenAiGateway extends AbstractOpenAiCompatibleGateway implements ModelProviderGateway {

    private static final Set<String> OPENAI_COMPATIBLE_PROVIDER_TYPES = Set.of(
            "openai",
            "deepseek",
            "aliyun",
            "zhipu",
            "minimax",
            "kimi",
            "siliconflow",
            "ollama",
            "self-hosted"
    );

    public OpenAiGateway(ObjectMapper objectMapper,
                         OpenAiProtocolMapper openAiProtocolMapper,
                         ThreadPoolTaskExecutor sseThreadPool) {
        super(objectMapper, openAiProtocolMapper, sseThreadPool);
    }

    @Override
    public boolean supports(String providerType) {
        return providerType != null && OPENAI_COMPATIBLE_PROVIDER_TYPES.contains(providerType.toLowerCase());
    }

    @Override
    public UnifiedChatResponse chat(Model model, UnifiedChatRequest request) {
        OpenAiChatCompletionRequest openAiRequest = openAiProtocolMapper.toOpenAiRequest(request);
        try (Response response = executePost(model, openAiRequest)) {
            validateResponse(response);
            if (response.body() == null) {
                throw tech.easyflow.ai.invoke.exception.ModelInvokeException.badGateway("上游模型响应为空");
            }
            OpenAiChatCompletionResponse responseBody = objectMapper.readValue(
                    response.body().string(),
                    OpenAiChatCompletionResponse.class
            );
            return openAiProtocolMapper.toUnifiedResponse(responseBody);
        } catch (tech.easyflow.ai.invoke.exception.ModelInvokeException e) {
            throw e;
        } catch (Exception e) {
            throw tech.easyflow.ai.invoke.exception.ModelInvokeException.badGateway("调用上游模型失败: " + e.getMessage());
        }
    }

    @Override
    public void chatStream(Model model, UnifiedChatRequest request, UnifiedChatChunkObserver observer) {
        OpenAiChatCompletionRequest openAiRequest = openAiProtocolMapper.toOpenAiRequest(request);
        Response response = executePost(model, openAiRequest);
        try {
            validateResponse(response);
            streamResponse(response, observer);
        } catch (RuntimeException | java.io.IOException e) {
            response.close();
            throw e instanceof RuntimeException ? (RuntimeException) e
                    : tech.easyflow.ai.invoke.exception.ModelInvokeException.badGateway("调用上游模型失败: " + e.getMessage());
        }
    }
}
