package com.easyagents.agent.runtime.agentscope;

import io.agentscope.core.message.*;
import io.agentscope.core.model.ChatResponse;
import io.agentscope.core.model.GenerateOptions;
import io.agentscope.core.model.Model;
import io.agentscope.core.model.ToolSchema;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * 面向 AutoContext 摘要调用的模型包装器。
 */
class AgentScopeAutoContextCompressionModel implements Model {

    private final Model delegate;

    /**
     * 创建 AutoContext 摘要模型包装器。
     *
     * @param delegate 实际调用的模型
     */
    AgentScopeAutoContextCompressionModel(Model delegate) {
        this.delegate = delegate;
    }

    /**
     * 对 AutoContext 压缩摘要输入做协议安全化后转发给实际模型。
     *
     * @param messages AgentScope 消息
     * @param tools 工具声明
     * @param options 生成参数
     * @return 模型响应流
     */
    @Override
    public Flux<ChatResponse> stream(List<Msg> messages, List<ToolSchema> tools, GenerateOptions options) {
        return delegate.stream(sanitizeCompressionMessages(messages), tools, options);
    }

    /**
     * 获取模型名称。
     *
     * @return 模型名称
     */
    @Override
    public String getModelName() {
        return delegate.getModelName();
    }

    /**
     * 将 AutoContext 摘要材料中的工具协议消息转换为普通文本消息。
     *
     * @param messages 原始摘要输入
     * @return 协议安全的摘要输入
     */
    List<Msg> sanitizeCompressionMessages(List<Msg> messages) {
        if (messages == null || messages.isEmpty()) {
            return messages;
        }
        List<Msg> sanitized = new ArrayList<>(messages.size());
        for (Msg message : messages) {
            if (message == null) {
                continue;
            }
            if (requiresPlainText(message)) {
                sanitized.add(toPlainTextUserMessage(message));
            } else {
                sanitized.add(message);
            }
        }
        return sanitized;
    }

    private boolean requiresPlainText(Msg message) {
        return message.getRole() == MsgRole.TOOL
                || message.hasContentBlocks(ToolUseBlock.class)
                || message.hasContentBlocks(ToolResultBlock.class);
    }

    private Msg toPlainTextUserMessage(Msg message) {
        return Msg.builder()
                .id(message.getId())
                .role(MsgRole.USER)
                .name("context")
                .content(TextBlock.builder().text(renderMessage(message)).build())
                .metadata(message.getMetadata())
                .timestamp(message.getTimestamp())
                .build();
    }

    private String renderMessage(Msg message) {
        StringBuilder builder = new StringBuilder();
        builder.append("Context message role=").append(message.getRole());
        if (message.getName() != null && !message.getName().isBlank()) {
            builder.append(", name=").append(message.getName());
        }
        builder.append('\n');
        for (ContentBlock block : message.getContent()) {
            appendBlock(builder, block);
        }
        return builder.toString().trim();
    }

    private void appendBlock(StringBuilder builder, ContentBlock block) {
        if (block instanceof TextBlock textBlock) {
            appendSection(builder, "text", textBlock.getText());
            return;
        }
        if (block instanceof ThinkingBlock thinkingBlock) {
            appendSection(builder, "thinking", thinkingBlock.getThinking());
            return;
        }
        if (block instanceof ToolUseBlock toolUseBlock) {
            appendSection(builder, "tool_use.id", toolUseBlock.getId());
            appendSection(builder, "tool_use.name", toolUseBlock.getName());
            appendSection(builder, "tool_use.input", renderMap(toolUseBlock.getInput()));
            return;
        }
        if (block instanceof ToolResultBlock toolResultBlock) {
            appendSection(builder, "tool_result.id", toolResultBlock.getId());
            appendSection(builder, "tool_result.name", toolResultBlock.getName());
            appendSection(builder, "tool_result.output", renderBlocks(toolResultBlock.getOutput()));
            return;
        }
        appendSection(builder, "content", String.valueOf(block));
    }

    private String renderBlocks(List<ContentBlock> blocks) {
        if (blocks == null || blocks.isEmpty()) {
            return "";
        }
        StringBuilder builder = new StringBuilder();
        for (ContentBlock block : blocks) {
            if (block instanceof TextBlock textBlock) {
                builder.append(textBlock.getText());
            } else {
                builder.append(String.valueOf(block));
            }
            builder.append('\n');
        }
        return builder.toString().trim();
    }

    private String renderMap(Map<String, Object> map) {
        return map == null ? "{}" : map.toString();
    }

    private void appendSection(StringBuilder builder, String label, String value) {
        if (value == null || value.isBlank()) {
            return;
        }
        builder.append(label).append(": ").append(value).append('\n');
    }
}
