package llama

import (
	"fmt"
	"strings"

	"git.db123.ir/db123/divacode/pkg/types"
)

type TemplateEngine struct {
	format string
}

func NewTemplateEngine(format string) *TemplateEngine {
	return &TemplateEngine{format: format}
}

func (te *TemplateEngine) Build(segments []types.PromptSegment) string {
	switch te.format {
	case "chatml":
		return te.buildChatML(segments)
	case "llama3":
		return te.buildLlama3(segments)
	default:
		return te.buildRaw(segments)
	}
}

func (te *TemplateEngine) StopTokens() []string {
	switch te.format {
	case "chatml":
		return []string{"<|im_end|>"}
	case "llama3":
		return []string{"<|eot_id|>"}
	default:
		return nil
	}
}

func (te *TemplateEngine) buildChatML(segments []types.PromptSegment) string {
	var b strings.Builder
	for _, seg := range segments {
		b.WriteString(fmt.Sprintf("<|im_start|>%s\n%s<|im_end|>\n", seg.Type, seg.Content))
	}
	b.WriteString("<|im_start|>assistant\n")
	return b.String()
}

func (te *TemplateEngine) buildLlama3(segments []types.PromptSegment) string {
	var b strings.Builder
	b.WriteString("<|begin_of_text|>")
	for _, seg := range segments {
		b.WriteString(fmt.Sprintf("<|start_header_id|>%s<|end_header_id|>\n\n%s<|eot_id|>", seg.Type, seg.Content))
	}
	b.WriteString("<|start_header_id|>assistant<|end_header_id|>\n\n")
	return b.String()
}

func (te *TemplateEngine) buildRaw(segments []types.PromptSegment) string {
	var b strings.Builder
	for _, seg := range segments {
		switch seg.Type {
		case types.SegSystem:
			b.WriteString("[System]\n" + seg.Content + "\n\n")
		case types.SegUser:
			b.WriteString("User: " + seg.Content + "\n")
		}
	}
	b.WriteString("You:")
	return b.String()
}