package com.pablotj.ia.chat.boot.application.prompt;

import java.util.ArrayList;
import java.util.List;
import org.thymeleaf.util.StringUtils;

public class PromptBuilder {

    private static final String END_TURN_SEPARATOR = "<|end_of_turn|>";
    private final String systemPrompt;
    private final List<String> turns = new ArrayList<>();
    public PromptBuilder(String systemPrompt) {
        this.systemPrompt = systemPrompt;
    }

    public void user(String message) {
        turns.add(this.formatMessage(MessageType.USER, message));
    }

    public void assistant(String message) {
        turns.add(this.formatMessage(MessageType.ASSISTANT, message));
    }

    public String build() {
        StringBuilder sb = new StringBuilder();
        sb.append(systemPrompt).append(END_TURN_SEPARATOR);
        for (String turn : turns) {
            sb.append(turn);
        }
        sb.append("GPT4 Correct Assistant:");
        return sb.toString();
    }

    private String formatMessage(MessageType messageType, String message) {
        return String.format("GPT4 Correct %s: %s %s", StringUtils.capitalize(messageType.name().toLowerCase()), message, END_TURN_SEPARATOR);
    }

    private enum MessageType {USER, ASSISTANT}
}