package com.example.demo;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.web.client.RestTemplate;
import org.json.JSONObject;
import org.json.JSONArray;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.junit.jupiter.api.Assertions.*;

class StructuredOutputsTest {

    private RestTemplate restTemplate;
    private HttpHeaders headers;
    private static final String API_KEY = "1";
    private static final String API_URL = "https://api.openai.com/v1/chat/completions";

    @BeforeEach
    void setUp() {
        restTemplate = new RestTemplate();
        headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("Authorization", "Bearer " + API_KEY);
    }

    @Test
    void testStructuredOutputs() {
        // Prepare the request body
        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", "gpt-4o-mini");
        requestBody.put("messages", List.of(
                Map.of("role", "system", "content", "You are a helpful math tutor. Guide the user through the solution step by step."),
                Map.of("role", "user", "content", "how can I solve 8x + 7 = -23")
        ));

        Map<String, Object> responseFormat = new HashMap<>();
        responseFormat.put("type", "json_schema");
        Map<String, Object> jsonSchema = new HashMap<>();
        jsonSchema.put("name", "math_reasoning");
        Map<String, Object> schema = new HashMap<>();
        schema.put("type", "object");
        Map<String, Object> properties = new HashMap<>();

        Map<String, Object> stepsProperty = new HashMap<>();
        stepsProperty.put("type", "array");
        Map<String, Object> items = new HashMap<>();
        items.put("type", "object");
        Map<String, Object> itemProperties = new HashMap<>();
        itemProperties.put("explanation", Map.of("type", "string"));
        itemProperties.put("output", Map.of("type", "string"));
        items.put("properties", itemProperties);
        items.put("required", List.of("explanation", "output"));
        items.put("additionalProperties", false);
        stepsProperty.put("items", items);

        properties.put("steps", stepsProperty);
        properties.put("final_answer", Map.of("type", "string"));

        schema.put("properties", properties);
        schema.put("required", List.of("steps", "final_answer"));
        schema.put("additionalProperties", false);

        jsonSchema.put("schema", schema);
        jsonSchema.put("strict", true);
        responseFormat.put("json_schema", jsonSchema);

        requestBody.put("response_format", responseFormat);

        // Make the API call
        HttpEntity<Map<String, Object>> request = new HttpEntity<>(requestBody, headers);
        String jsonResponse = restTemplate.postForObject(API_URL, request, String.class);

        // Parse and validate the response
        JSONObject response = new JSONObject(jsonResponse);
        JSONObject content = new JSONObject(response.getJSONArray("choices")
                .getJSONObject(0)
                .getJSONObject("message")
                .getString("content"));
        System.out.println("Response content: " + content.toString(2));

        // Assertions
        assertTrue(content.has("steps"));
        assertTrue(content.has("final_answer"));

        JSONArray stepsArray = content.getJSONArray("steps");
        assertEquals(5, stepsArray.length());

        for (int i = 0; i < stepsArray.length(); i++) {
            JSONObject step = stepsArray.getJSONObject(i);
            assertTrue(step.has("explanation"));
            assertTrue(step.has("output"));
            assertFalse(step.getString("explanation").isEmpty());
            assertFalse(step.getString("output").isEmpty());
        }

        // 수정된 부분: 최종 답변 검증
        String finalAnswer = content.getString("final_answer");
        assertTrue(finalAnswer.equals("x = -15 / 4") || finalAnswer.equals("x = -3.75"),
                "Final answer should be either 'x = -15 / 4' or 'x = -3.75'");
    }

출력을 json으로 받을수있다니

이얼마나 멋진일 아닙니까?

보자마자 바로 테스트 해봤지요 물론 코드작성은 인공지능이.. ㅋ

자바로 테스트 코드 작성해서 실행해보니 잘되는군요

json으로 출력이 된다니 이거 활용방안이 무궁무진합니다요