package com.zy.ai.service;

import com.alibaba.fastjson.JSON;
import com.zy.ai.entity.ChatCompletionRequest;
import com.zy.ai.entity.ChatCompletionResponse;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.http.client.reactive.ReactorClientHttpConnector;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.client.WebClient;
import io.netty.channel.ChannelOption;
import reactor.netty.http.client.HttpClient;
import reactor.core.publisher.Flux;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;

@Service
public class LlmSpringAiClientService {

    @Value("${llm.connect-timeout-ms:10000}")
    private int connectTimeoutMs;

    @Value("${llm.read-timeout-ms:120000}")
    private int readTimeoutMs;

    public CompletionCallResult callCompletion(String baseUrl, String apiKey, ChatCompletionRequest req) {
        OpenAiApi api = buildOpenAiApi(baseUrl, apiKey);
        OpenAiApi.ChatCompletionRequest springReq = buildSpringAiRequest(req, false);
        ResponseEntity<OpenAiApi.ChatCompletion> entity = api.chatCompletionEntity(springReq);
        OpenAiApi.ChatCompletion body = entity.getBody();
        ChatCompletionResponse legacy = toLegacyResponse(body);
        return new CompletionCallResult(entity.getStatusCode().value(),
                serializeResponsePayload(body, legacy),
                legacy);
    }

    public Flux<String> streamCompletion(String baseUrl, String apiKey, ChatCompletionRequest req) {
        OpenAiApi api = buildOpenAiApi(baseUrl, apiKey);
        OpenAiApi.ChatCompletionRequest springReq = buildSpringAiRequest(req, true);
        return api.chatCompletionStream(springReq)
                .flatMapIterable(chunk -> chunk == null || chunk.choices() == null
                        ? List.<OpenAiApi.ChatCompletionChunk.ChunkChoice>of()
                        : chunk.choices())
                .map(OpenAiApi.ChatCompletionChunk.ChunkChoice::delta)
                .filter(delta -> delta != null)
                .handle((delta, sink) -> {
                    String text = extractSpringAiContent(delta);
                    if (text != null && !text.isEmpty()) {
                        sink.next(text);
                    }
                });
    }

    public Integer statusCodeOf(Throwable ex) {
        if (ex == null || ex.getMessage() == null) {
            return null;
        }
        String text = ex.getMessage().trim();
        int idx = text.indexOf(" - ");
        String prefix = idx >= 0 ? text.substring(0, idx).trim() : text;
        try {
            return Integer.parseInt(prefix);
        } catch (NumberFormatException ignore) {
            return null;
        }
    }

    public String responseBodyOf(Throwable ex, int maxLen) {
        if (ex == null || ex.getMessage() == null) {
            return null;
        }
        String text = ex.getMessage().trim();
        int idx = text.indexOf(" - ");
        String body = idx >= 0 ? text.substring(idx + 3).trim() : text;
        if (body.length() > maxLen) {
            return body.substring(0, maxLen);
        }
        return body;
    }

    private OpenAiApi buildOpenAiApi(String baseUrl, String apiKey) {
        HttpClient httpClient = HttpClient.create()
                .option(ChannelOption.CONNECT_TIMEOUT_MILLIS, connectTimeoutMs)
                .responseTimeout(Duration.ofMillis(readTimeoutMs));
        WebClient.Builder webClientBuilder = WebClient.builder()
                .clientConnector(new ReactorClientHttpConnector(httpClient));
        RestClient.Builder restClientBuilder = RestClient.builder()
                .requestFactory(buildRequestFactory());
        return OpenAiApi.builder()
                .baseUrl(baseUrl)
                .apiKey(apiKey)
                .restClientBuilder(restClientBuilder)
                .webClientBuilder(webClientBuilder)
                .build();
    }

    private SimpleClientHttpRequestFactory buildRequestFactory() {
        SimpleClientHttpRequestFactory factory = new SimpleClientHttpRequestFactory();
        factory.setConnectTimeout(connectTimeoutMs);
        factory.setReadTimeout(readTimeoutMs);
        return factory;
    }

    private OpenAiApi.ChatCompletionRequest buildSpringAiRequest(ChatCompletionRequest req, boolean stream) {
        HashMap<String, Object> extraBody = new HashMap<>();
        if (req != null && req.getThinking() != null) {
            HashMap<String, Object> thinking = new HashMap<>();
            thinking.put("type", req.getThinking().getType() != null ? req.getThinking().getType() : "enable");
            extraBody.put("thinking", thinking);
        }
        return new OpenAiApi.ChatCompletionRequest(
                toSpringAiMessages(req == null ? null : req.getMessages()),
                req == null ? null : req.getModel(),
                null,
                null,
                null,
                null,
                null,
                null,
                req == null ? null : req.getMax_tokens(),
                null,
                1,
                null,
                null,
                null,
                null,
                null,
                null,
                null,
                stream,
                stream ? OpenAiApi.ChatCompletionRequest.StreamOptions.INCLUDE_USAGE : null,
                req == null ? null : req.getTemperature(),
                null,
                toSpringAiTools(req == null ? null : req.getTools()),
                toSpringAiToolChoice(req == null ? null : req.getTool_choice()),
                null,
                null,
                null,
                null,
                null,
                null,
                null,
                extraBody.isEmpty() ? null : extraBody
        );
    }

    private List<OpenAiApi.ChatCompletionMessage> toSpringAiMessages(List<ChatCompletionRequest.Message> messages) {
        ArrayList<OpenAiApi.ChatCompletionMessage> result = new ArrayList<>();
        if (messages == null) {
            return result;
        }
        for (ChatCompletionRequest.Message message : messages) {
            if (message == null) {
                continue;
            }
            result.add(new OpenAiApi.ChatCompletionMessage(
                    message.getContent(),
                    toSpringAiRole(message.getRole()),
                    message.getName(),
                    message.getTool_call_id(),
                    toSpringAiToolCalls(message.getTool_calls()),
                    null,
                    null,
                    null,
                    null
            ));
        }
        return result;
    }

    private List<OpenAiApi.ChatCompletionMessage.ToolCall> toSpringAiToolCalls(List<ChatCompletionRequest.ToolCall> toolCalls) {
        if (toolCalls == null || toolCalls.isEmpty()) {
            return null;
        }
        ArrayList<OpenAiApi.ChatCompletionMessage.ToolCall> result = new ArrayList<>();
        for (ChatCompletionRequest.ToolCall toolCall : toolCalls) {
            if (toolCall == null) {
                continue;
            }
            ChatCompletionRequest.Function function = toolCall.getFunction();
            OpenAiApi.ChatCompletionMessage.ChatCompletionFunction springFunction =
                    new OpenAiApi.ChatCompletionMessage.ChatCompletionFunction(
                            function == null ? null : function.getName(),
                            function == null ? null : function.getArguments()
                    );
            result.add(new OpenAiApi.ChatCompletionMessage.ToolCall(
                    toolCall.getId(),
                    toolCall.getType(),
                    springFunction
            ));
        }
        return result.isEmpty() ? null : result;
    }

    private OpenAiApi.ChatCompletionMessage.Role toSpringAiRole(String role) {
        if (role == null) {
            return OpenAiApi.ChatCompletionMessage.Role.USER;
        }
        switch (role.trim().toLowerCase(Locale.ROOT)) {
            case "system":
                return OpenAiApi.ChatCompletionMessage.Role.SYSTEM;
            case "assistant":
                return OpenAiApi.ChatCompletionMessage.Role.ASSISTANT;
            case "tool":
                return OpenAiApi.ChatCompletionMessage.Role.TOOL;
            default:
                return OpenAiApi.ChatCompletionMessage.Role.USER;
        }
    }

    private List<OpenAiApi.FunctionTool> toSpringAiTools(List<Object> tools) {
        if (tools == null || tools.isEmpty()) {
            return null;
        }
        ArrayList<OpenAiApi.FunctionTool> result = new ArrayList<>();
        for (Object tool : tools) {
            OpenAiApi.FunctionTool mapped = toSpringAiTool(tool);
            if (mapped != null) {
                result.add(mapped);
            }
        }
        return result.isEmpty() ? null : result;
    }

    @SuppressWarnings("unchecked")
    private OpenAiApi.FunctionTool toSpringAiTool(Object tool) {
        if (tool instanceof OpenAiApi.FunctionTool) {
            return (OpenAiApi.FunctionTool) tool;
        }
        if (!(tool instanceof Map)) {
            return null;
        }
        Map<String, Object> toolMap = (Map<String, Object>) tool;
        Object functionObj = toolMap.get("function");
        if (!(functionObj instanceof Map)) {
            return null;
        }
        Map<String, Object> functionMap = (Map<String, Object>) functionObj;
        Object name = functionMap.get("name");
        if (name == null) {
            return null;
        }
        Object description = functionMap.get("description");
        Object parameters = functionMap.get("parameters");
        Map<String, Object> parametersMap;
        if (parameters instanceof Map) {
            parametersMap = new HashMap<>((Map<String, Object>) parameters);
        } else {
            parametersMap = Collections.emptyMap();
        }
        OpenAiApi.FunctionTool.Function function = new OpenAiApi.FunctionTool.Function(
                String.valueOf(description == null ? "" : description),
                String.valueOf(name),
                parametersMap,
                null
        );
        return new OpenAiApi.FunctionTool(OpenAiApi.FunctionTool.Type.FUNCTION, function);
    }

    @SuppressWarnings("unchecked")
    private Object toSpringAiToolChoice(Object toolChoice) {
        if (toolChoice == null) {
            return null;
        }
        if (toolChoice instanceof String) {
            String text = ((String) toolChoice).trim();
            if (text.isEmpty()) {
                return null;
            }
            if ("auto".equalsIgnoreCase(text) || "none".equalsIgnoreCase(text)) {
                return text.toLowerCase(Locale.ROOT);
            }
            return OpenAiApi.ChatCompletionRequest.ToolChoiceBuilder.function(text);
        }
        if (toolChoice instanceof Map) {
            Map<String, Object> choiceMap = (Map<String, Object>) toolChoice;
            Object type = choiceMap.get("type");
            Object function = choiceMap.get("function");
            if ("function".equals(type) && function instanceof Map) {
                Object name = ((Map<String, Object>) function).get("name");
                if (name != null) {
                    return OpenAiApi.ChatCompletionRequest.ToolChoiceBuilder.function(String.valueOf(name));
                }
            }
        }
        return toolChoice;
    }

    private ChatCompletionResponse toLegacyResponse(OpenAiApi.ChatCompletion completion) {
        if (completion == null) {
            return null;
        }
        ChatCompletionResponse response = new ChatCompletionResponse();
        response.setId(completion.id());
        response.setCreated(completion.created());
        response.setObjectName(completion.object());
        if (completion.usage() != null) {
            ChatCompletionResponse.Usage usage = new ChatCompletionResponse.Usage();
            usage.setPromptTokens(completion.usage().promptTokens());
            usage.setCompletionTokens(completion.usage().completionTokens());
            usage.setTotalTokens(completion.usage().totalTokens());
            response.setUsage(usage);
        }
        if (completion.choices() != null) {
            ArrayList<ChatCompletionResponse.Choice> choices = new ArrayList<>();
            for (OpenAiApi.ChatCompletion.Choice choice : completion.choices()) {
                ChatCompletionResponse.Choice item = new ChatCompletionResponse.Choice();
                item.setIndex(choice.index());
                if (choice.finishReason() != null) {
                    item.setFinishReason(choice.finishReason().name().toLowerCase(Locale.ROOT));
                }
                item.setMessage(toLegacyMessage(choice.message()));
                choices.add(item);
            }
            response.setChoices(choices);
        }
        return response;
    }

    private String serializeResponsePayload(OpenAiApi.ChatCompletion body, ChatCompletionResponse legacy) {
        if (legacy != null) {
            String content = firstMessageContent(legacy);
            if (content != null && !content.trim().isEmpty()) {
                return content;
            }
            String legacyJson = JSON.toJSONString(legacy);
            if (legacyJson != null && !legacyJson.trim().isEmpty() && !"{}".equals(legacyJson.trim())) {
                return legacyJson;
            }
        }
        if (body != null) {
            String raw = body.toString();
            if (raw != null && !raw.trim().isEmpty()) {
                return raw;
            }
        }
        return null;
    }

    private String firstMessageContent(ChatCompletionResponse response) {
        if (response == null || response.getChoices() == null || response.getChoices().isEmpty()) {
            return null;
        }
        ChatCompletionResponse.Choice choice = response.getChoices().get(0);
        if (choice == null || choice.getMessage() == null) {
            return null;
        }
        return choice.getMessage().getContent();
    }

    private ChatCompletionRequest.Message toLegacyMessage(OpenAiApi.ChatCompletionMessage message) {
        if (message == null) {
            return null;
        }
        ChatCompletionRequest.Message result = new ChatCompletionRequest.Message();
        result.setContent(extractSpringAiContent(message));
        if (message.role() != null) {
            result.setRole(message.role().name().toLowerCase(Locale.ROOT));
        }
        result.setName(message.name());
        result.setTool_call_id(message.toolCallId());
        result.setTool_calls(toLegacyToolCalls(message.toolCalls()));
        return result;
    }

    private List<ChatCompletionRequest.ToolCall> toLegacyToolCalls(List<OpenAiApi.ChatCompletionMessage.ToolCall> toolCalls) {
        if (toolCalls == null || toolCalls.isEmpty()) {
            return null;
        }
        ArrayList<ChatCompletionRequest.ToolCall> result = new ArrayList<>();
        for (OpenAiApi.ChatCompletionMessage.ToolCall toolCall : toolCalls) {
            if (toolCall == null) {
                continue;
            }
            ChatCompletionRequest.ToolCall legacy = new ChatCompletionRequest.ToolCall();
            legacy.setId(toolCall.id());
            legacy.setType(toolCall.type());
            if (toolCall.function() != null) {
                ChatCompletionRequest.Function function = new ChatCompletionRequest.Function();
                function.setName(toolCall.function().name());
                function.setArguments(toolCall.function().arguments());
                legacy.setFunction(function);
            }
            result.add(legacy);
        }
        return result.isEmpty() ? null : result;
    }

    private String extractSpringAiContent(OpenAiApi.ChatCompletionMessage message) {
        if (message == null || message.rawContent() == null) {
            return null;
        }
        Object content = message.rawContent();
        if (content instanceof String) {
            return (String) content;
        }
        if (content instanceof List) {
            try {
                @SuppressWarnings("unchecked")
                List<OpenAiApi.ChatCompletionMessage.MediaContent> media =
                        (List<OpenAiApi.ChatCompletionMessage.MediaContent>) content;
                return OpenAiApi.getTextContent(media);
            } catch (ClassCastException ignore) {
            }
        }
        return String.valueOf(content);
    }

    public static class CompletionCallResult {
        private final int statusCode;
        private final String payload;
        private final ChatCompletionResponse response;

        public CompletionCallResult(int statusCode, String payload, ChatCompletionResponse response) {
            this.statusCode = statusCode;
            this.payload = payload;
            this.response = response;
        }

        public int getStatusCode() {
            return statusCode;
        }

        public String getPayload() {
            return payload;
        }

        public ChatCompletionResponse getResponse() {
            return response;
        }
    }
}