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.mockito.ArgumentCaptor;
import org.springframework.ai.chat.client.ChatClient;

import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
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#extractWithRetry} 双层重试单元测试。
 *
 * <p>覆盖：</p>
 * <ul>
 *   <li>外层（运行重试）：LLM 调用异常或空响应时重新执行完整推理</li>
 *   <li>内层（格式重试）：中文键名、字段缺失、类型不匹配时仅重组 JSON，不重跑业务逻辑</li>
 *   <li>内层耗尽升级到外层</li>
 *   <li>重组 prompt 内容验证（不含原 userPrompt、列出缺失键名）</li>
 * </ul>
 */
class StructuredOutputHelperDualLoopTest {

    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) {
        // 默认 required=true（匹配修订后 strict 模式语义：只对 required 字段做键名校验）
        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);
        }
    }

    /**
     * 配置 mock ChatClient 按 prompt 内容返回不同响应：
     * - 含 "重新组织" → 重组调用
     * - 其他 → 完整推理调用
     */
    private void mockResponsesByPrompt(String fullResponse, String reformatResponse) {
        when(callResponse.content()).thenAnswer(inv -> {
            // 通过 spec.user(...) 的最近一次调用判断
            // 这里简化：使用顺序——第一次调用返回 fullResponse，后续返回 reformatResponse
            return null; // 由下方 answer 替换
        });
    }

    /**
     * 按调用顺序返回响应。
     */
    private void mockResponsesInOrder(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：首次 OK 不重试 ==========

    @Test
    void firstCallOk_noRetry() {
        JsonNode decl = outputsDecl("status|number", "message|string");
        mockResponsesInOrder("{\"status\":200,\"message\":\"ok\"}");

        RetryPolicy policy = RetryPolicy.defaultPolicy();
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetry(client, "do work", null, decl, policy);

        assertThat(result.errorMessage).isNull();
        assertThat(result.data).containsEntry("status", 200).containsEntry("message", "ok");
        assertThat(result.attempts).isEqualTo(1);
        verify(spec, times(1)).call();
    }

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

    @Test
    void chineseKeys_triggerFormatRetry() {
        JsonNode decl = outputsDecl("status|number", "message|string");
        // 第 1 次：返回中文键名；第 2 次：内层重组，返回英文键名
        mockResponsesInOrder(
                "{\"状态\":200,\"信息\":\"ok\"}",
                "{\"status\":200,\"message\":\"ok\"}"
        );

        RetryPolicy policy = RetryPolicy.defaultPolicy();
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetry(client, "do work", null, decl, policy);

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

    // ========== 用例 3：LLM 异常 → 外层重试 ==========

    @Test
    void runtimeException_triggersRuntimeRetry() {
        JsonNode decl = outputsDecl("status|number");
        // 第 1 次抛异常；第 2 次返回正确 JSON
        when(callResponse.content())
                .thenThrow(new RuntimeException("rate limit"))
                .thenReturn("{\"status\":200}");

        RetryPolicy policy = RetryPolicy.defaultPolicy();
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetry(client, "do work", null, decl, policy);

        assertThat(result.errorMessage).isNull();
        assertThat(result.data).containsEntry("status", 200);
        verify(spec, times(2)).call();
    }

    // ========== 用例 4：内层重试耗尽 → 升级到外层 ==========

    @Test
    void formatRetryExhausted_upgradesToRuntime() {
        JsonNode decl = outputsDecl("status|number", "message|string");
        // 第 1 次外层：返回中文；第 2 次内层重组：仍是中文（formatMaxRetries=2，耗尽）
        // 第 3 次外层：返回英文 → 成功
        mockResponsesInOrder(
                "{\"状态\":200,\"信息\":\"ok\"}",     // 外层 1
                "{\"状态\":200,\"信息\":\"ok\"}",     // 内层重组 1（仍中文）
                "{\"status\":200,\"message\":\"ok\"}" // 外层 2（成功）
        );

        // formatMaxRetries=2：内层最多尝试 2 次（首次解析 + 1 次重组）
        RetryPolicy policy = new RetryPolicy(2, 2, true, true, true, true);
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetry(client, "do work", null, decl, policy);

        assertThat(result.errorMessage).isNull();
        assertThat(result.data).containsEntry("status", 200).containsEntry("message", "ok");
        // 1（外层1）+ 1（内层重组1）+ 1（外层2）= 3 次调用
        verify(spec, times(3)).call();
    }

    // ========== 用例 5：所有重试都耗尽 → 返回失败 ==========

    @Test
    void allRetriesExhausted_returnsFail() {
        JsonNode decl = outputsDecl("status|number");
        // runtimeMaxRetries=2, formatMaxRetries=2 → 每次外层触发 1 次完整 + 1 次重组 = 2 次调用
        // 2 个外层 → 总 4 次 LLM 调用
        mockResponsesInOrder("{\"状态\":200}", "{\"状态\":200}", "{\"状态\":200}", "{\"状态\":200}");

        RetryPolicy policy = new RetryPolicy(2, 2, true, true, true, true);
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetry(client, "do work", null, decl, policy);

        assertThat(result.errorMessage).isNotNull();
        assertThat(result.errorMessage).contains("status");
        assertThat(result.data).isEmpty();
        // 2 次外层完整推理 + 2 次内层重组 = 4 次
        verify(spec, times(4)).call();
    }

    // ========== 用例 6：重组 prompt 不包含原 userPrompt ==========

    @Test
    void reformatPrompt_doesNotContainUserPrompt() {
        JsonNode decl = outputsDecl("status|number");
        mockResponsesInOrder(
                "{\"状态\":200}",          // 外层返回中文
                "{\"status\":200}"          // 内层重组成功
        );

        String userPrompt = "BUSINESS_SECRET_PROMPT_DO_NOT_LEAK";
        RetryPolicy policy = RetryPolicy.defaultPolicy();
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetry(client, userPrompt, null, decl, policy);

        assertThat(result.errorMessage).isNull();

        // 捕获所有 prompt 调用
        ArgumentCaptor<String> promptCaptor = ArgumentCaptor.forClass(String.class);
        verify(spec, times(2)).user(promptCaptor.capture());

        // 第 1 次：完整 prompt 应包含 userPrompt
        assertThat(promptCaptor.getAllValues().get(0)).contains(userPrompt);
        // 第 2 次：重组 prompt 不应包含 userPrompt
        assertThat(promptCaptor.getAllValues().get(1)).doesNotContain(userPrompt);
    }

    // ========== 用例 7：重组 prompt 明确列出缺失键名 ==========

    @Test
    void reformatPrompt_listsAllMissingKeys() {
        JsonNode decl = outputsDecl("status|number", "message|string", "data|object");
        // 返回全中文键 → 所有英文声明都缺失
        mockResponsesInOrder(
                "{\"状态\":200,\"信息\":\"ok\",\"数据\":{}}",
                "{\"status\":200,\"message\":\"ok\",\"data\":{}}"
        );

        RetryPolicy policy = RetryPolicy.defaultPolicy();
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetry(client, "work", null, decl, policy);

        assertThat(result.errorMessage).isNull();

        ArgumentCaptor<String> promptCaptor = ArgumentCaptor.forClass(String.class);
        verify(spec, times(2)).user(promptCaptor.capture());

        // 重组 prompt 必须列出所有缺失的英文键名
        String reformatPrompt = promptCaptor.getAllValues().get(1);
        assertThat(reformatPrompt).contains("status")
                .contains("message")
                .contains("data");
        assertThat(reformatPrompt).contains("重新组织");
    }

    // ========== 用例 8：retryOnKeyNameMismatch=false 时不重组 ==========

    @Test
    void keyNameMismatchDisabled_noReformatRetry() {
        JsonNode decl = outputsDecl("status|number");
        mockResponsesInOrder("{\"状态\":200}");

        // 关闭键名不匹配重试
        RetryPolicy policy = new RetryPolicy(2, 2, true, true, true, false);
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetry(client, "work", null, decl, policy);

        assertThat(result.errorMessage).isNotNull();
        assertThat(result.errorMessage).contains("status");
        // 仅 1 次调用，不重组
        verify(spec, times(1)).call();
    }

    // ========== 用例 9：空响应 → 外层重试 ==========

    @Test
    void emptyResponse_triggersRuntimeRetry() {
        JsonNode decl = outputsDecl("status|number");
        // 第 1 次空；第 2 次正常
        when(callResponse.content())
                .thenReturn("")
                .thenReturn("{\"status\":200}");

        RetryPolicy policy = RetryPolicy.defaultPolicy();
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetry(client, "work", null, decl, policy);

        assertThat(result.errorMessage).isNull();
        assertThat(result.data).containsEntry("status", 200);
        verify(spec, times(2)).call();
    }

    // ========== 用例 10：类型不匹配 → 内层重组 ==========

    @Test
    void typeMismatch_triggersFormatRetry() {
        JsonNode decl = outputsDecl("status|number");
        // 第 1 次：status 是字符串 "abc"，无法转 number；第 2 次重组为数字
        mockResponsesInOrder(
                "{\"status\":\"abc\"}",
                "{\"status\":200}"
        );

        RetryPolicy policy = RetryPolicy.defaultPolicy();
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetry(client, "work", null, decl, policy);

        assertThat(result.errorMessage).isNull();
        assertThat(result.data).containsEntry("status", 200);
        verify(spec, times(2)).call();
    }

    // ========== 用例 11：parse 失败（非 JSON） → 内层重组 ==========

    @Test
    void parseFailure_triggersFormatRetry() {
        JsonNode decl = outputsDecl("status|number");
        mockResponsesInOrder(
                "this is not json at all",
                "{\"status\":200}"
        );

        RetryPolicy policy = RetryPolicy.defaultPolicy();
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetry(client, "work", null, decl, policy);

        assertThat(result.errorMessage).isNull();
        assertThat(result.data).containsEntry("status", 200);
        verify(spec, times(2)).call();
    }

    // ========== 用例 12：自定义 systemPrompt 被传递 ==========

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

        String systemPrompt = "you are a JSON generator";
        RetryPolicy policy = RetryPolicy.defaultPolicy();
        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetry(client, "work", systemPrompt, decl, policy);

        assertThat(result.errorMessage).isNull();
        verify(spec, times(1)).system(systemPrompt);
    }

    // ========== 用例 13：null client / null decl 直接失败 ==========

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

        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetry(null, "work", null, decl, policy);

        assertThat(result.errorMessage).contains("ChatClient");
        assertThat(result.attempts).isEqualTo(0);
    }

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

        StructuredOutputHelper.ExtractionResult result =
                StructuredOutputHelper.extractWithRetry(client, "work", null, null, policy);

        assertThat(result.errorMessage).contains("outputs");
        assertThat(result.attempts).isEqualTo(0);
    }
}
