Prebuilt Components

/**
 * Agent state specialized for message collections
 * @param <T> Type of messages stored in the state
 */
class MessagesState<T> extends AgentState {
    /**
     * Default schema with appender channel for messages
     */
    static final Map<String, Channel<?>> SCHEMA;

    /**
     * Creates MessagesState with initial data
     * @param initData Initial state data including messages
     */
    MessagesState(Map<String, Object> initData);

    /**
     * Retrieves the list of messages
     * @return List of messages in the state
     * @throws RuntimeException if messages key not found
     */
    List<T> messages();

    /**
     * Gets the most recent message
     * @return Optional containing last message if present
     */
    Optional<T> lastMessage();

    /**
     * Gets message at position N from the end
     * @param n Position from end (0 = last, 1 = second to last, etc.)
     * @return Optional containing message at position if exists
     */
    Optional<T> lastMinus(int n);
}

// String-based messages
class StringMessagesState extends MessagesState<String> {
    public StringMessagesState(Map<String, Object> initData) {
        super(initData);
    }
}

// Create initial state with messages
List<String> initialMessages = List.of("Hello", "How can I help?");
Map<String, Object> initData = Map.of("messages", new ArrayList<>(initialMessages));
StringMessagesState state = new StringMessagesState(initData);

// Access messages
List<String> allMessages = state.messages();
Optional<String> lastMsg = state.lastMessage();
Optional<String> secondToLast = state.lastMinus(1);

System.out.println("All messages: " + allMessages);
System.out.println("Last message: " + lastMsg.orElse("None"));
System.out.println("Previous message: " + secondToLast.orElse("None"));

// Custom message type
class ChatMessage {
    private final String content;
    private final String sender;
    private final long timestamp;

    public ChatMessage(String content, String sender) {
        this.content = content;
        this.sender = sender;
        this.timestamp = System.currentTimeMillis();
    }

    // getters...
    public String getContent() { return content; }
    public String getSender() { return sender; }
    public long getTimestamp() { return timestamp; }
}

class ChatMessagesState extends MessagesState<ChatMessage> {
    public ChatMessagesState(Map<String, Object> initData) {
        super(initData);
    }
}

// Use with custom message type
List<ChatMessage> chatMessages = List.of(
    new ChatMessage("Hello there!", "user"),
    new ChatMessage("Hi! How can I assist you?", "assistant")
);

ChatMessagesState chatState = new ChatMessagesState(
    Map.of("messages", new ArrayList<>(chatMessages))
);

Optional<ChatMessage> lastChatMsg = chatState.lastMessage();
if (lastChatMsg.isPresent()) {
    System.out.println("Last from: " + lastChatMsg.get().getSender());
    System.out.println("Content: " + lastChatMsg.get().getContent());
}

/**
 * StateGraph specialized for MessagesState
 * @param <T> Type of messages in the state
 */
class MessagesStateGraph<T> extends StateGraph<MessagesState<T>> {
    /**
     * Creates MessagesStateGraph with default constructor factory
     */
    MessagesStateGraph();

    /**
     * Creates MessagesStateGraph with custom state serializer
     * @param stateSerializer Serializer for MessagesState instances
     */
    MessagesStateGraph(StateSerializer<MessagesState<T>> stateSerializer);
}

// Simple messages graph with default serialization
MessagesStateGraph<String> simpleMessagesGraph = new MessagesStateGraph<>();

// Add nodes that work with messages
simpleMessagesGraph.addNode("receive_input", (state) -> {
    String userInput = state.<String>value("user_input").orElse("");
    // Messages are automatically appended due to MessagesState.SCHEMA
    return CompletableFuture.completedFuture(Map.of("messages", List.of(userInput)));
});

simpleMessagesGraph.addNode("generate_response", (state) -> {
    List<String> messages = state.messages();
    String lastMessage = state.lastMessage().orElse("");

    // Simple echo response
    String response = "Echo: " + lastMessage;
    return CompletableFuture.completedFuture(Map.of("messages", List.of(response)));
});

simpleMessagesGraph.addNode("display_conversation", (state) -> {
    List<String> messages = state.messages();
    System.out.println("Conversation history:");
    for (int i = 0; i < messages.size(); i++) {
        System.out.println((i + 1) + ". " + messages.get(i));
    }
    return CompletableFuture.completedFuture(Map.of("displayed", true));
});

// Set up flow
simpleMessagesGraph.addEdge(StateGraph.START, "receive_input");
simpleMessagesGraph.addEdge("receive_input", "generate_response");
simpleMessagesGraph.addEdge("generate_response", "display_conversation");
simpleMessagesGraph.addEdge("display_conversation", StateGraph.END);

// Compile and run
CompiledGraph<MessagesState<String>> app = simpleMessagesGraph.compile();
Optional<MessagesState<String>> result = app.invoke(Map.of("user_input", "Hello AI!"));

if (result.isPresent()) {
    List<String> finalMessages = result.get().messages();
    System.out.println("Final conversation: " + finalMessages);
}

// Advanced messages graph with custom serialization
StateSerializer<MessagesState<ChatMessage>> chatSerializer =
    new JacksonStateSerializer<>(ChatMessagesState::new);

MessagesStateGraph<ChatMessage> chatGraph = new MessagesStateGraph<>(chatSerializer);

// Add sophisticated chat nodes
chatGraph.addNode("process_user_message", (state) -> {
    String userInput = state.<String>value("user_input").orElse("");
    ChatMessage userMsg = new ChatMessage(userInput, "user");

    return CompletableFuture.completedFuture(Map.of(
        "messages", List.of(userMsg),
        "last_user_input", userInput
    ));
});

chatGraph.addNode("generate_ai_response", (state) -> {
    String lastUserInput = state.<String>value("last_user_input").orElse("");

    // Simulate AI processing based on conversation history
    List<ChatMessage> history = state.messages();
    String context = "Previous messages: " + history.size();
    String response = generateAIResponse(lastUserInput, context);

    ChatMessage aiMsg = new ChatMessage(response, "assistant");

    return CompletableFuture.completedFuture(Map.of(
        "messages", List.of(aiMsg),
        "response_generated", true
    ));
});

// Chat routing based on message analysis
chatGraph.addNode("analyze_intent", (state, config) -> {
    Optional<ChatMessage> lastMsg = state.lastMessage();

    return CompletableFuture.supplyAsync(() -> {
        if (lastMsg.isPresent() && "user".equals(lastMsg.get().getSender())) {
            String content = lastMsg.get().getContent().toLowerCase();

            String intent;
            if (content.contains("goodbye") || content.contains("bye")) {
                intent = "farewell";
            } else if (content.contains("help") || content.contains("?")) {
                intent = "help_request";
            } else {
                intent = "general_chat";
            }

            return new Command(intent, Map.of("detected_intent", intent));
        }

        return new Command("general_chat", Map.of());
    });
}, Map.of(
    "farewell", "farewell_handler",
    "help_request", "help_handler",
    "general_chat", "general_chat_handler"
));

private String generateAIResponse(String input, String context) {
    // Simulate AI response generation
    return "AI response to: " + input + " (Context: " + context + ")";
}

// Message filtering and transformation
AsyncNodeAction<MessagesState<ChatMessage>> messageFilter = (state) -> {
    List<ChatMessage> messages = state.messages();

    // Filter out system messages older than 1 hour
    long oneHourAgo = System.currentTimeMillis() - 3600000;
    List<ChatMessage> filtered = messages.stream()
        .filter(msg -> !"system".equals(msg.getSender()) || msg.getTimestamp() > oneHourAgo)
        .collect(Collectors.toList());

    return CompletableFuture.completedFuture(Map.of(
        "messages", filtered,
        "filtered_count", messages.size() - filtered.size()
    ));
};

// Message summarization
AsyncNodeAction<MessagesState<String>> messageSummarizer = (state) -> {
    List<String> messages = state.messages();

    return CompletableFuture.supplyAsync(() -> {
        if (messages.size() > 10) {
            // Summarize older messages, keep recent ones
            List<String> recent = messages.subList(messages.size() - 5, messages.size());
            String summary = "Summary of " + (messages.size() - 5) + " earlier messages: " +
                messages.subList(0, messages.size() - 5).stream()
                    .collect(Collectors.joining(", "));

            List<String> condensed = new ArrayList<>();
            condensed.add(summary);
            condensed.addAll(recent);

            return Map.of("messages", condensed, "summarized", true);
        }

        return Map.of("summarized", false);
    });
};

// Message validation
AsyncNodeAction<MessagesState<String>> messageValidator = (state) -> {
    Optional<String> lastMsg = state.lastMessage();

    return CompletableFuture.completedFuture(
        lastMsg.map(msg -> {
            boolean isValid = msg.length() > 0 && msg.length() <= 1000 &&
                             !msg.trim().isEmpty();

            Map<String, Object> updates = new HashMap<>();
            updates.put("last_message_valid", isValid);

            if (!isValid) {
                updates.put("messages", List.of("Please provide a valid message (1-1000 characters)"));
            }

            return updates;
        }).orElse(Map.of("last_message_valid", false))
    );
};

// Conversation state tracking
class ConversationState extends MessagesState<ChatMessage> {
    public ConversationState(Map<String, Object> initData) {
        super(initData);
    }

    public Optional<String> getCurrentTopic() {
        return this.value("current_topic");
    }

    public Optional<Integer> getMessageCount() {
        return Optional.of(messages().size());
    }

    public boolean hasRecentActivity() {
        return lastMessage()
            .map(msg -> System.currentTimeMillis() - msg.getTimestamp() < 300000) // 5 min
            .orElse(false);
    }
}

// Conversation management graph
MessagesStateGraph<ChatMessage> conversationGraph = new MessagesStateGraph<ChatMessage>() {
    @Override
    protected ConversationState getStateFactory() {
        return ConversationState::new;
    }
};

// Topic tracking node
conversationGraph.addNode("track_topic", (state) -> {
    Optional<ChatMessage> lastMsg = state.lastMessage();

    return CompletableFuture.supplyAsync(() -> {
        String topic = "general";

        if (lastMsg.isPresent()) {
            String content = lastMsg.get().getContent().toLowerCase();
            if (content.contains("weather")) topic = "weather";
            else if (content.contains("food") || content.contains("recipe")) topic = "cooking";
            else if (content.contains("help") || content.contains("support")) topic = "support";
        }

        return Map.of(
            "current_topic", topic,
            "topic_changed", !topic.equals(state.<String>value("current_topic").orElse("general"))
        );
    });
});

// Context-aware response generation
conversationGraph.addNode("contextual_response", (state) -> {
    List<ChatMessage> messages = state.messages();
    String currentTopic = state.<String>value("current_topic").orElse("general");

    return CompletableFuture.supplyAsync(() -> {
        // Generate response based on conversation context
        String response = generateContextualResponse(messages, currentTopic);
        ChatMessage aiResponse = new ChatMessage(response, "assistant");

        return Map.of(
            "messages", List.of(aiResponse),
            "response_context", currentTopic
        );
    });
});

// Conversation cleanup
conversationGraph.addNode("cleanup_old_messages", (state) -> {
    List<ChatMessage> messages = state.messages();

    return CompletableFuture.completedFuture(
        messages.size() > 50 ?
            Map.of("messages", messages.subList(messages.size() - 30, messages.size())) :
            Map.of("cleanup_performed", false)
    );
});

private String generateContextualResponse(List<ChatMessage> messages, String topic) {
    // Context-aware response generation logic
    return "Response for topic '" + topic + "' based on " + messages.size() + " messages";
}

// Hybrid state combining messages with other data
class HybridState extends MessagesState<String> {
    public HybridState(Map<String, Object> initData) {
        super(initData);
    }

    public Optional<Integer> getScore() {
        return this.value("score");
    }

    public Optional<String> getUserId() {
        return this.value("user_id");
    }
}

// Custom schema combining messages with other channels
Map<String, Channel<?>> hybridSchema = new HashMap<>(MessagesState.SCHEMA);
hybridSchema.put("score", (key, current, newVal) ->
    Math.max((Integer) (current != null ? current : 0), (Integer) newVal));
hybridSchema.put("analytics", Channels.appender(ArrayList::new));

// Create graph with hybrid schema
StateGraph<HybridState> hybridGraph = new StateGraph<>(hybridSchema, HybridState::new);

hybridGraph.addNode("process_message_with_scoring", (state) -> {
    Optional<String> lastMsg = state.lastMessage();
    int currentScore = state.getScore().orElse(0);

    return CompletableFuture.supplyAsync(() -> {
        Map<String, Object> updates = new HashMap<>();

        if (lastMsg.isPresent()) {
            String msg = lastMsg.get();
            updates.put("messages", List.of("Processed: " + msg));

            // Calculate score based on message
            int msgScore = msg.length() > 10 ? 10 : 5;
            updates.put("score", msgScore);

            // Add analytics event
            updates.put("analytics", List.of(Map.of(
                "event", "message_processed",
                "score_delta", msgScore,
                "timestamp", System.currentTimeMillis()
            )));
        }

        return updates;
    });
});

// Compile and use hybrid graph
CompiledGraph<HybridState> hybridApp = hybridGraph.compile();
Optional<HybridState> hybridResult = hybridApp.invoke(Map.of(
    "messages", List.of("Hello hybrid world!"),
    "user_id", "user123",
    "score", 0
));

if (hybridResult.isPresent()) {
    HybridState finalState = hybridResult.get();
    System.out.println("Messages: " + finalState.messages());
    System.out.println("Final score: " + finalState.getScore().orElse(0));
    System.out.println("User: " + finalState.getUserId().orElse("unknown"));
}