package com.huimv.employment.integration.kb;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.huimv.employment.common.exception.BizException;
import com.huimv.employment.common.exception.ErrorCode;
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.util.StringUtils;
import org.springframework.web.client.HttpStatusCodeException;
import org.springframework.web.client.RestTemplate;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;

/**
 * 将智能体 {@code /api/version}、{@code /api/console/chat} 转发至上游网关。
 */
@Service
public class KbOpenAiProxyServiceImpl implements KbOpenAiProxyService {

    private static final String HEADER_AGENT_ID = "X-Agent-Id";
    private static final String PATH_VERSION = "/api/version";
    private static final String PATH_CHAT = "/api/console/chat";

    private final KbApiProperties properties;
    private final RestTemplate restTemplate;

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

    @Override
    public ResponseEntity<byte[]> getVersion() {
        ensureEnabled();
        try {
            ResponseEntity<byte[]> upstream = restTemplate.exchange(
                    properties.baseUrl() + PATH_VERSION,
                    HttpMethod.GET,
                    new HttpEntity<>(upstreamHeaders(null, false)),
                    byte[].class);
            return forward(upstream);
        } catch (HttpStatusCodeException ex) {
            return forwardError(ex);
        }
    }

    @Override
    public void streamConsoleChat(byte[] requestBody, String agentId, OutputStream output) throws IOException {
        streamConsoleChatInternal(requestBody, agentId, output, null);
    }

    @Override
    public ConsoleChatCollectResult streamConsoleChatCollecting(byte[] requestBody, String agentId, OutputStream output,
                                                                ObjectMapper objectMapper) throws IOException {
        ConsoleChatSseAggregator aggregator = new ConsoleChatSseAggregator(objectMapper);
        streamConsoleChatInternal(requestBody, agentId, output, aggregator);
        aggregator.finish();
        return new ConsoleChatCollectResult(
                aggregator.getAssistantText(),
                aggregator.isCompleted(),
                aggregator.takeHeldTerminalLine(),
                aggregator.getMcpDraftEvidence(),
                aggregator.getLastMcpDraftId());
    }

    private void streamConsoleChatInternal(byte[] requestBody, String agentId, OutputStream output,
                                           ConsoleChatSseAggregator aggregator) throws IOException {
        ensureEnabled();
        validateRequestBody(requestBody);
        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);
            applyUpstreamAuth(conn, agentId);
            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) {
                    writeUpstreamError(output, input);
                } else if (aggregator != null) {
                    // 会话落库：暂缓 [DONE]，便于业务层把卡片追加到最后消息后再结束
                    SseStreamSupport.pipeCollectingDeferDone(input, output, aggregator);
                } else {
                    SseStreamSupport.pipeWithFlush(input, output, null);
                }
            }
        } finally {
            if (input != null) {
                try {
                    input.close();
                } catch (IOException ignored) {
                    // ignore
                }
            }
            if (conn != null) {
                conn.disconnect();
            }
        }
    }

    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 applyUpstreamAuth(HttpURLConnection conn, String agentId) {
        conn.setRequestProperty(HEADER_AGENT_ID, resolveAgentId(agentId));
        if (properties.hasUpstreamAuth()) {
            conn.setRequestProperty(properties.getAuthHeaderName(), properties.authHeaderValue());
        }
    }

    private void writeUpstreamError(OutputStream output, InputStream input) throws IOException {
        byte[] errBytes = StreamUtils.copyToByteArray(input);
        if (errBytes.length == 0) {
            return;
        }
        String errText = new String(errBytes, StandardCharsets.UTF_8).trim();
        if (errText.startsWith("data:") || errText.contains("\n\n")) {
            output.write(errBytes);
            output.flush();
        } else {
            SseStreamSupport.writeDataEvent(output, errText);
        }
    }

    private String resolveAgentId(String agentId) {
        if (StringUtils.hasText(agentId)) {
            return agentId.trim();
        }
        return StringUtils.hasText(properties.getDefaultAgentId())
                ? properties.getDefaultAgentId().trim()
                : "default";
    }

    private void validateRequestBody(byte[] requestBody) {
        if (requestBody == null || requestBody.length == 0) {
            throw new BizException(ErrorCode.BAD_REQUEST, "请求体不能为空");
        }
    }

    private void ensureEnabled() {
        if (!properties.isEnabled()) {
            throw new BizException(ErrorCode.AI_SERVICE_DISABLED);
        }
    }

    private HttpHeaders upstreamHeaders(String agentId, boolean requireAgentId) {
        HttpHeaders headers = new HttpHeaders();
        properties.applyAuthHeader(headers);
        if (requireAgentId) {
            headers.add(HEADER_AGENT_ID, resolveAgentId(agentId));
        }
        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());
    }
}
