package com.ruoyi.web.kb.service.impl;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;

import com.ruoyi.web.kb.support.KbApiProperties;
import com.ruoyi.web.kb.support.SseStreamSupport;
import com.ruoyi.web.kb.service.KbOpenAiProxyService;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.util.StreamUtils;
import org.springframework.web.client.HttpStatusCodeException;
import org.springframework.web.client.RestTemplate;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.ruoyi.common.exception.ServiceException;
import com.ruoyi.common.utils.StringUtils;

/**
 * 将 OpenAI 兼容对话/模型列表接口转发至上游大模型网关，配置与 {@link KnowledgeBaseClientImpl} 一致。
 */
@Service
public class KbOpenAiProxyServiceImpl implements KbOpenAiProxyService
{
    private static final String PATH_MODELS = "/v1/models";

    private static final String PATH_CHAT = "/v1/chat/completions";

    private final KbApiProperties properties;

    private final RestTemplate restTemplate;

    private final ObjectMapper objectMapper;

    public KbOpenAiProxyServiceImpl(KbApiProperties properties,
            @Qualifier("kbRestTemplate") RestTemplate restTemplate,
            ObjectMapper objectMapper)
    {
        this.properties = properties;
        this.restTemplate = restTemplate;
        this.objectMapper = objectMapper;
    }

    @Override
    public ResponseEntity<byte[]> listModels()
    {
        ensureReady();
        try
        {
            ResponseEntity<byte[]> upstream = restTemplate.exchange(
                    properties.baseUrl() + PATH_MODELS,
                    HttpMethod.GET,
                    new HttpEntity<>(upstreamAuthHeaders()),
                    byte[].class);
            return forward(upstream);
        }
        catch (HttpStatusCodeException ex)
        {
            return forwardError(ex);
        }
    }

    @Override
    public ResponseEntity<byte[]> chatCompletions(byte[] requestBody)
    {
        ensureReady();
        if (requestBody == null || requestBody.length == 0)
        {
            throw new ServiceException("请求体不能为空");
        }
        try
        {
            HttpHeaders headers = upstreamAuthHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            ResponseEntity<byte[]> upstream = restTemplate.exchange(
                    properties.baseUrl() + PATH_CHAT,
                    HttpMethod.POST,
                    new HttpEntity<>(requestBody, headers),
                    byte[].class);
            return forward(upstream);
        }
        catch (HttpStatusCodeException ex)
        {
            return forwardError(ex);
        }
    }

    @Override
    public void streamChatCompletions(byte[] requestBody, OutputStream output) throws IOException
    {
        ensureReady();
        if (requestBody == null || requestBody.length == 0)
        {
            throw new ServiceException("请求体不能为空");
        }
        HttpURLConnection conn = null;
        InputStream input = null;
        try
        {
            conn = openUpstreamConnection(requestBody.length);
            conn.setRequestProperty(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE);
            conn.setRequestProperty(HttpHeaders.ACCEPT, MediaType.TEXT_EVENT_STREAM_VALUE);
            conn.setRequestProperty(properties.getAuthHeaderName(), properties.authHeaderValue());
            conn.setDoOutput(true);
            try (OutputStream upstreamOut = conn.getOutputStream())
            {
                upstreamOut.write(requestBody);
            }
            int status = conn.getResponseCode();
            input = status >= 400 ? conn.getErrorStream() : conn.getInputStream();
            if (input != null)
            {
                if (status >= 400)
                {
                    byte[] errBytes = StreamUtils.copyToByteArray(input);
                    if (errBytes.length > 0)
                    {
                        String errText = new String(errBytes, java.nio.charset.StandardCharsets.UTF_8).trim();
                        if (errText.startsWith("data:") || errText.contains("\n\n"))
                        {
                            output.write(errBytes);
                            output.flush();
                        }
                        else
                        {
                            SseStreamSupport.writeDataEvent(output, errText);
                        }
                    }
                }
                else
                {
                    SseStreamSupport.pipeWithFlush(input, output);
                }
            }
        }
        finally
        {
            if (input != null)
            {
                try
                {
                    input.close();
                }
                catch (IOException ignored)
                {
                    /* ignore */
                }
            }
            if (conn != null)
            {
                conn.disconnect();
            }
        }
    }

    @Override
    public boolean isStreamRequest(byte[] requestBody)
    {
        if (requestBody == null || requestBody.length == 0)
        {
            return false;
        }
        try
        {
            JsonNode root = objectMapper.readTree(requestBody);
            JsonNode stream = root.get("stream");
            return stream != null && stream.asBoolean(false);
        }
        catch (Exception ex)
        {
            return false;
        }
    }

    private HttpURLConnection openUpstreamConnection(int contentLength) throws IOException
    {
        URL url = new URL(properties.baseUrl() + PATH_CHAT);
        HttpURLConnection conn = (HttpURLConnection) url.openConnection();
        conn.setRequestMethod("POST");
        conn.setConnectTimeout(properties.getConnectTimeoutMs());
        conn.setReadTimeout(properties.getStreamReadTimeoutMs());
        conn.setFixedLengthStreamingMode(contentLength);
        return conn;
    }

    private void ensureReady()
    {
        if (!properties.isEnabled())
        {
            throw new ServiceException("大模型网关已关闭（ruoyi.kb.enabled=false）");
        }
        if (StringUtils.isEmpty(properties.getApiKey()))
        {
            throw new ServiceException("未配置大模型访问密钥 ruoyi.kb.api-key");
        }
    }

    private HttpHeaders upstreamAuthHeaders()
    {
        HttpHeaders headers = new HttpHeaders();
        headers.add(properties.getAuthHeaderName(), properties.authHeaderValue());
        return headers;
    }

    private ResponseEntity<byte[]> forward(ResponseEntity<byte[]> upstream)
    {
        HttpHeaders headers = new HttpHeaders();
        MediaType contentType = upstream.getHeaders().getContentType();
        if (contentType != null)
        {
            headers.setContentType(contentType);
        }
        else
        {
            headers.setContentType(MediaType.APPLICATION_JSON);
        }
        return new ResponseEntity<>(upstream.getBody(), headers, upstream.getStatusCode());
    }

    private ResponseEntity<byte[]> forwardError(HttpStatusCodeException ex)
    {
        HttpHeaders headers = new HttpHeaders();
        MediaType contentType = ex.getResponseHeaders() != null ? ex.getResponseHeaders().getContentType() : null;
        if (contentType != null)
        {
            headers.setContentType(contentType);
        }
        else
        {
            headers.setContentType(MediaType.APPLICATION_JSON);
        }
        byte[] body = ex.getResponseBodyAsByteArray();
        return new ResponseEntity<>(body, headers, ex.getStatusCode());
    }
}
