package com.agent.management.engine;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.client.ChatClient;

import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

/**
 * {@link StructuredOutputHelper#extractWithRetryFromText} 单元测试。
 *
 * <p>覆盖 Hermes / Agent 系列接入入口的核心场景：</p>
 * <ul>
 *   <li>原始文本直接解析成功（不需重组）</li>
 *   <li>原始文本需内层重组（用 reformatClient 整理为 JSON）</li>
 *   <li>内层耗尽 → supplier 重新生成 rawText（外层重试）</li>
 *   <li>supplier 为 null 时不做外层重试</li>
 *   <li>所有重试耗尽返回失败</li>
 * </ul>
 */
class StructuredOutputFromTextTest {

    private static final ObjectMapper MAPPER = new ObjectMapper();

    private ChatClient client;
    private ChatClient.ChatClientRequestSpec spec;
    private ChatClient.CallResponseSpec callResponse;

    @BeforeEach
    void setUp() {
        client = mock(ChatClient.class);
        spec = mock(ChatClient.ChatClientRequestSpec.class);
        callResponse = mock(ChatClient.CallResponseSpec.class);

        when(client.prompt()).thenReturn(spec);
        when(spec.user(anyString())).thenReturn(spec);
        when(spec.system(anyString())).thenReturn(spec);
        when(spec.call()).thenReturn(callResponse);
    }

    private static JsonNode outputsDecl(String... fields) {
        StringBuilder sb = new StringBuilder("[");
        for (int i = 0; i < fields.length; i++) {
            if (i > 0) sb.append(",");
            String[] parts = fields[i].split("\\|");
            sb.append("{\"name\":\"").append(parts[0]).append("\",")
              .append("\"type\":\"").append(parts[1]).append("\",")
              .append("\"required\":true");
            sb.append("}");
        }
        sb.append("]");
        try {
            return MAPPER.readTree(sb.toString());
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private void mockReformatResponses(String... responses) {
        AtomicInteger idx = new AtomicInteger(0);
        when(callResponse.content()).thenAnswer(inv -> {
            int i = idx.getAndIncrement();
            return i < responses.length ? responses[i] : responses[responses.length - 1];
        });
    }

    // ========== 用例 1：rawText 直接是有效 JSON ==========

    @Test
    void rawTextValidJson_noReformatNeeded() {
        JsonNode decl = outputsDecl("status|number", "message|string");
        // reformatClient 不会被调用
        RetryPolicy policy = RetryPolicy.defaultPolicy();

        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetryFromText(
                        "{\"status\":200,\"message\":\"ok\"}",
                        decl, policy, client, null);

        assertThat(result.errorMessage).isNull();
        assertThat(result.data).containsEntry("status", 200).containsEntry("message", "ok");
        // 不调用 reformatClient
        verify(spec, times(0)).call();
    }

    // ========== 用例 2：rawText 中文键名 → 内层重组成功 ==========

    @Test
    void chineseKeys_inRawText_triggerReformat() {
        JsonNode decl = outputsDecl("status|number");
        mockReformatResponses("{\"status\":200}");

        RetryPolicy policy = RetryPolicy.defaultPolicy();
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetryFromText(
                        "{\"状态\":200}", decl, policy, client, null);

        assertThat(result.errorMessage).isNull();
        assertThat(result.data).containsEntry("status", 200);
        // 1 次内层重组调用
        verify(spec, times(1)).call();
    }

    // ========== 用例 3：内层耗尽 → supplier 重新生成 → 成功 ==========

    @Test
    void reformatExhausted_supplierProvidesNewText_succeeds() {
        JsonNode decl = outputsDecl("status|number");
        // 第 1 次外层：rawText 是中文，内层重组 1 次仍中文 → 内层耗尽（formatMaxRetries=2 时最多调 1 次重组）
        // 第 2 次外层：supplier 返回英文 JSON → 直接解析成功，不再重组
        mockReformatResponses("{\"状态\":200}");

        AtomicInteger supplierCalls = new AtomicInteger(0);
        Supplier<String> supplier = () -> {
            supplierCalls.incrementAndGet();
            return "{\"status\":200}";
        };

        // runtimeMaxRetries=2, formatMaxRetries=2
        RetryPolicy policy = new RetryPolicy(2, 2, true, true, true, true);
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetryFromText(
                        "{\"状态\":200}", decl, policy, client, supplier);

        assertThat(result.errorMessage).isNull();
        assertThat(result.data).containsEntry("status", 200);
        assertThat(supplierCalls.get()).isEqualTo(1);
        // 第 1 次外层触发 1 次重组；第 2 次外层 rawText 直接解析成功不重组
        verify(spec, times(1)).call();
    }

    // ========== 用例 4：supplier 为 null 时直接失败 ==========

    @Test
    void supplierNull_innerExhausted_fails() {
        JsonNode decl = outputsDecl("status|number");
        // 内层重组 2 次都返回中文 → 耗尽
        mockReformatResponses(
                "{\"状态\":200}",
                "{\"状态\":200}"
        );

        RetryPolicy policy = new RetryPolicy(2, 2, true, true, true, true);
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetryFromText(
                        "{\"状态\":200}", decl, policy, client, null);

        assertThat(result.errorMessage).isNotNull();
        // supplier 为 null，不进入外层重试，重组调用 = formatMaxRetries-1 = 1
        verify(spec, times(1)).call();
    }

    // ========== 用例 5：rawText 空字符串 → 失败 ==========

    @Test
    void emptyRawText_returnsFail() {
        JsonNode decl = outputsDecl("status|number");
        RetryPolicy policy = RetryPolicy.defaultPolicy();

        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetryFromText(
                        "", decl, policy, client, null);

        assertThat(result.errorMessage).contains("原始响应为空");
        assertThat(result.attempts).isEqualTo(0);
    }

    // ========== 用例 6：null outputsDecl → 失败 ==========

    @Test
    void nullOutputsDecl_returnsFail() {
        RetryPolicy policy = RetryPolicy.defaultPolicy();

        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetryFromText(
                        "anything", null, policy, client, null);

        assertThat(result.errorMessage).contains("outputs");
    }

    // ========== 用例 7：null reformatClient → 失败 ==========

    @Test
    void nullReformatClient_returnsFail() {
        JsonNode decl = outputsDecl("status|number");
        RetryPolicy policy = RetryPolicy.defaultPolicy();

        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetryFromText(
                        "{\"status\":200}", decl, policy, null, null);

        assertThat(result.errorMessage).contains("reformatClient");
    }

    // ========== 用例 8：supplier 异常 → 升级失败 ==========

    @Test
    void supplierThrows_exceptionPropagatedAsError() {
        JsonNode decl = outputsDecl("status|number");
        // 内层重组耗尽
        mockReformatResponses(
                "{\"状态\":200}",
                "{\"状态\":200}"
        );

        AtomicInteger supplierCalls = new AtomicInteger(0);
        Supplier<String> supplier = () -> {
            supplierCalls.incrementAndGet();
            throw new RuntimeException("hermes down");
        };

        RetryPolicy policy = new RetryPolicy(2, 2, true, true, true, true);
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetryFromText(
                        "{\"状态\":200}", decl, policy, client, supplier);

        assertThat(result.errorMessage).contains("hermes down");
        assertThat(supplierCalls.get()).isEqualTo(1);
    }

    // ========== 用例 9：rawText 是非 JSON 自由文本 → 内层重组 ==========

    @Test
    void freeTextRawText_reformatByLlm() {
        JsonNode decl = outputsDecl("summary|string");
        // Hermes finalText 是自由文本，内层 LLM 把它整理成 JSON
        mockReformatResponses("{\"summary\":\"这是总结\"}");

        RetryPolicy policy = RetryPolicy.defaultPolicy();
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetryFromText(
                        "经过分析，结论是：这是总结。",
                        decl, policy, client, null);

        assertThat(result.errorMessage).isNull();
        assertThat(result.data).containsEntry("summary", "这是总结");
        verify(spec, times(1)).call();
    }
}
