package llm

import "github.com/tmc/langchaingo/llms"

type MessageBuilder struct {
	messages []llms.MessageContent
}

func NewMessageBuilder() *MessageBuilder {
	return &MessageBuilder{
		messages: make([]llms.MessageContent, 0),
	}
}

func (b *MessageBuilder) WithSystemPrompt(prompt string) *MessageBuilder {
	b.messages = append(b.messages, llms.TextParts(llms.ChatMessageTypeSystem, prompt))
	return b
}

func (b *MessageBuilder) WithMessages(msgs []llms.MessageContent) *MessageBuilder {
	b.messages = append(b.messages, msgs...)
	return b
}

func (b *MessageBuilder) AppendAIMessage(content string, toolCalls []llms.ToolCall) *MessageBuilder {
	parts := []llms.ContentPart{}
	if content != "" {
		parts = append(parts, llms.TextPart(content))
	}
	for _, tc := range toolCalls {
		parts = append(parts, tc)
	}
	b.messages = append(b.messages, llms.MessageContent{
		Role:  llms.ChatMessageTypeAI,
		Parts: parts,
	})
	return b
}

func (b *MessageBuilder) AppendToolResult(toolCallID, name, result string) *MessageBuilder {
	b.messages = append(b.messages, llms.MessageContent{
		Role: llms.ChatMessageTypeTool,
		Parts: []llms.ContentPart{
			llms.ToolCallResponse{
				ToolCallID: toolCallID,
				Name:       name,
				Content:    result,
			},
		},
	})
	return b
}

func (b *MessageBuilder) Build() []llms.MessageContent {
	result := make([]llms.MessageContent, len(b.messages))
	copy(result, b.messages)
	return result
}