package com.infoepoch.pms.agent;

import com.infoepoch.pms.agent.properties.AiZhiYuProperties;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.time.Duration;

@SpringBootTest
public class LLMChatTest {

    private static final Logger log = LoggerFactory.getLogger(LLMChatTest.class);

    private static final String PROMPT_TEXT = "Say this is a test!";
    private static final Duration HTTP_TIMEOUT = Duration.ofSeconds(60);

    private final ChatModel pmsChatModel;
    private final AiZhiYuProperties properties;

    @Autowired
    public LLMChatTest(ChatModel pmsChatModel, AiZhiYuProperties properties) {
        this.pmsChatModel = pmsChatModel;
        this.properties = properties;
    }

    @Test
    public void chatTest() {
        logRuntimeContext();
        runSpringAiCall();
        runRawHttpComparison();
    }

    private void runSpringAiCall() {
        log.info("[spring-ai] starting request");
        try {
            Prompt prompt = new Prompt(PROMPT_TEXT,
                    OpenAiChatOptions.builder()
                            .model(properties.getModel())
                            .temperature(properties.getTemperature())
                            .maxTokens(properties.getMaxTokens())
                            .build());

            ChatResponse chatResponse = pmsChatModel.call(prompt);
            String output = chatResponse == null || chatResponse.getResult() == null
                    || chatResponse.getResult().getOutput() == null
                    ? null
                    : chatResponse.getResult().getOutput().getText();
            log.info("[spring-ai] response_text={}", output);
        } catch (Exception exception) {
            log.error("[spring-ai] request failed", exception);
        }
    }

    private void runRawHttpComparison() {
        String requestBody = buildRawRequestBody();
        String requestUrl = buildChatCompletionsUrl();
        log.info("[raw-http] starting request url={}", requestUrl);
        log.info("[raw-http] request_headers={{Authorization=Bearer ****, Content-Type=application/json}}");
        log.info("[raw-http] request_body_bytes={}", requestBody.getBytes(StandardCharsets.UTF_8).length);
        log.info("[raw-http] request_body={}", requestBody);

        try {
            HttpRequest request = HttpRequest.newBuilder()
                    .uri(URI.create(requestUrl))
                    .timeout(HTTP_TIMEOUT)
                    .header("Authorization", "Bearer " + resolveApiKey())
                    .header("Content-Type", "application/json")
                    .POST(HttpRequest.BodyPublishers.ofString(requestBody, StandardCharsets.UTF_8))
                    .build();

            HttpResponse<String> response = HttpClient.newBuilder()
                    .connectTimeout(HTTP_TIMEOUT)
                    .build()
                    .send(request, HttpResponse.BodyHandlers.ofString(StandardCharsets.UTF_8));

            log.info("[raw-http] response_status={}", response.statusCode());
            log.info("[raw-http] response_headers={}", response.headers().map());
            log.info("[raw-http] response_body={}", response.body());
        } catch (Exception exception) {
            log.error("[raw-http] request failed", exception);
        }
    }

    private void logRuntimeContext() {
        log.info("[diagnostic] java_version={}", System.getProperty("java.version"));
        log.info("[diagnostic] java_vendor={}", System.getProperty("java.vendor"));
        log.info("[diagnostic] os_name={}", System.getProperty("os.name"));
        log.info("[diagnostic] base_url={}", properties.getBaseUrl());
        log.info("[diagnostic] model={}", properties.getModel());
        log.info("[diagnostic] prompt={}", PROMPT_TEXT);
    }

    private String buildRawRequestBody() {
        return """
                {
                  \"model\": \"%s\",
                  \"messages\": [
                    {
                      \"role\": \"user\",
                      \"content\": \"%s\"
                    }
                  ],
                  \"temperature\": %s
                }
                """.formatted(properties.getModel(), escapeJson(PROMPT_TEXT), properties.getTemperature());
    }

    private String buildChatCompletionsUrl() {
        String baseUrl = properties.getBaseUrl();
        if (baseUrl.endsWith("/")) {
            return baseUrl + "v1/chat/completions";
        }
        return baseUrl + "/v1/chat/completions";
    }

    private String resolveApiKey() {
        return properties.getApiKey() == null ? "EMPTY" : properties.getApiKey();
    }

    private String escapeJson(String value) {
        return value
                .replace("\\", "\\\\")
                .replace("\"", "\\\"");
    }
}
