package features.ai.team.level;

import demo.ai.llm.LlmUtil;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import org.noear.solon.ai.agent.Agent;
import org.noear.solon.ai.agent.AgentSession;
import org.noear.solon.ai.agent.session.InMemoryAgentSession;
import org.noear.solon.ai.agent.team.TeamAgent;
import org.noear.solon.ai.agent.team.TeamTrace;
import org.noear.solon.ai.chat.ChatModel;
import org.noear.solon.ai.chat.message.AssistantMessage;
import org.noear.solon.ai.chat.message.ChatMessage;
import org.noear.solon.ai.chat.prompt.Prompt;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.stream.Collectors;

/**
 * TeamAgent 递归与循环协作测试
 * * <p>场景 1: 测试嵌套团队（Team in Team），验证 Trace 轨迹在父子团队间的正确记录。</p>
 * <p>场景 2: 测试反馈循环（Feedback Loop），验证当审核不通过时，流程能重新回到研发节点。</p>
 */
public class TeamAgentRecursiveTest {
    private static final Logger log = LoggerFactory.getLogger(TeamAgentRecursiveTest.class);

    /**
     * 测试：嵌套团队协作
     * 验证：父团队（项目组）调用子团队（研发组）时，各层级的 Trace 均能正确生成且不触发死循环。
     */
    @Test
    public void testNestedTeam() throws Throwable {
        ChatModel chatModel = LlmUtil.getChatModel();

        // 1. 底层团队：研发小组 (dev_team)
        // 调整：使用 role().instruction() 风格定义职责
        TeamAgent devTeam = TeamAgent.of(chatModel)
                .name("dev_team")
                .role("研发小组")
                .instruction("输入需求，直接给出代码实现。完成后必须回复：[FINISH] 研发已完成")
                .agentAdd(createSimpleAgent("Coder", "负责写代码"))
                .maxTurns(2) // 严格限制子团队迭代次数
                .build();

        // 2. 顶层团队：项目管理组 (project_team)
        // 调整：弃用 systemPrompt，改用新的 role().instruction()
        TeamAgent projectTeam = TeamAgent.of(chatModel)
                .name("project_team")
                .role("严谨的项目主管")
                .instruction("1. 先指派 Analyst 进行分析；" +
                        "2. 拿到分析结果后指派 dev_team 执行；" +
                        "3. 当 dev_team 回复完成后，直接输出 [FINISH]。")
                .agentAdd(createSimpleAgent("Analyst", "需求分析师"))
                .agentAdd(devTeam) // 嵌套子团队
                .maxTurns(5)
                .build();

        // 可视化结构：打印图定义的 YAML
        System.out.println("--- Project Team Graph ---\n" + projectTeam.getGraph().toYaml() + "\n---");

        // 3. 使用 AgentSession 替代 FlowContext
        AgentSession session = InMemoryAgentSession.of("sn_recursive_2026");

        log.info(">>> 开始嵌套团队调用测试...");
        // 核心：在 Prompt 中明确要求一次性处理，调整为新的调用风格
        String promptText = "请 Java 程序员帮我写一个 Hello World。完成后直接结束。";
        projectTeam.prompt(Prompt.of(promptText)).session(session).call().getContent();

        // 4. 结果验证：从 session 的快照中提取 Trace
        TeamTrace rootTrace = session.getContext().getAs("__project_team");
        TeamTrace subTrace = session.getContext().getAs("__dev_team");

        if (rootTrace != null) {
            log.info("父团队执行路径: {}", String.join(" -> ",
                    rootTrace.getRecords().stream().map(s -> s.getSource()).toArray(String[]::new)));
        }

        Assertions.assertNotNull(rootTrace, "父团队 Trace 记录丢失");
        Assertions.assertTrue(rootTrace.getTurnCount() < 5, "触发了非预期的高频迭代，可能存在逻辑死循环");

        log.info("子团队是否存在轨迹: {}", (subTrace != null));
    }

    /**
     * 测试：反馈修正循环
     * 验证：当 Reviewer 给出打回意见时，TeamAgent 能自动调度回 dev_team 进行修复。
     */
    @Test
    public void testFeedbackLoop() throws Throwable {
        ChatModel chatModel = LlmUtil.getChatModel();

        // 1. 简单的开发子团队
        TeamAgent devTeam = TeamAgent.of(chatModel).name("dev_team")
                .role("代码实现小组")
                .instruction("负责根据主管要求编写或修复代码实现。")
                .feedbackMode(false)
                .agentAdd(createSimpleAgent("Coder", "程序员"))
                .build();

        // 3. 带有审核逻辑的顶层团队
        // 调整：合并 role 定义并使用 instruction 设置协作红线
        TeamAgent projectTeam = TeamAgent.of(chatModel).name("quality_project")
                .role("严谨的代码项目主管")
                .instruction("## 协作红线：\n" +
                        "1. 所有的开发工作（dev_team）完成后，必须指派 Reviewer 进行审核。\n" +
                        "2. 严禁直接结束任务，除非 Reviewer 给出确认通过的结论。\n" +
                        "3. 如果 Reviewer 指出问题，你必须重新指派 dev_team 修复并说明审核意见。")
                .feedbackMode(false)
                .agentAdd(devTeam)
                .agentAdd(new Agent() {
                    private int reviewCount = 0;
                    @Override public String name() { return "Reviewer"; }
                    @Override public String role() { return "代码审核员"; }

                    @Override
                    public AssistantMessage call(Prompt prompt, AgentSession session) {
                        if (reviewCount++ == 0) {
                            return ChatMessage.ofAssistant("代码发现安全漏洞，请 dev_team 重新修复！[FINISH]");
                        }
                        return ChatMessage.ofAssistant("审核通过，表现完美。[FINISH]");
                    }
                })
                .maxTurns(10)
                .build();

        // 4. 执行任务
        AgentSession session = InMemoryAgentSession.of("sn_feedback_loop_2026");
        // 调整：改为新的链式调用风格
        String result = projectTeam.prompt(Prompt.of("请使用 Java 编写一个简单的登录 Controller 接口。"))
                .session(session)
                .call()
                .getContent();

        System.out.println("=====最终输出=====");
        System.out.println(result);

        // 5. 获取执行轨迹
        TeamTrace rootTrace = session.getContext().getAs(projectTeam.getConfig().getTraceKey());
        Assertions.assertNotNull(rootTrace, "执行轨迹丢失");

        // 打印历史日志供分析
        System.out.println("--- Collaboration History Log ---");
        System.out.println(rootTrace.getFormattedHistory());
        System.out.println("--- End History Log ---");

        // 验证逻辑...
        List<String> routeHistory = rootTrace.getRecords().stream()
                .filter(TeamTrace.TeamRecord::isAgent)
                .map(s -> s.getSource())
                .collect(Collectors.toList());

        long devTeamCalls = routeHistory.stream()
                .filter(name -> "dev_team".equalsIgnoreCase(name)).count();

        log.info("反馈循环中 dev_team 被激活次数: {}", devTeamCalls);
        Assertions.assertTrue(devTeamCalls >= 2, "当审核未通过时，应重新指派 dev_team 修复");
        Assertions.assertTrue(routeHistory.size() >= 4, "完整的反馈循环至少需要 4 个步骤");
        Assertions.assertEquals("dev_team", routeHistory.get(0), "首位执行者应该是 dev_team");

        boolean hasApproval = rootTrace.getRecords().stream()
                .anyMatch(s -> "Reviewer".equals(s.getSource()) && s.getContent().contains("表现完美"));

        Assertions.assertTrue(hasApproval, "最终结果应包含 Reviewer 的通过确认");
        Assertions.assertNotNull(rootTrace.getFinalAnswer(), "任务结束时应有最终答案产出");
    }

    private static Agent createSimpleAgent(String name, String desc) {
        return new Agent() {
            @Override public String name() { return name; }
            @Override public String role() { return desc; }
            @Override
            public AssistantMessage call(Prompt prompt, AgentSession session) {
                return ChatMessage.ofAssistant("[Result from " + name + "]: 任务处理完毕。 [FINISH]");
            }
        };
    }
}