package prompt

import (
	"bytes"
	"fmt"
	"gitee.com/kzangv/gsf-ai-agent/utils"
	"regexp"
	"text/template"
)

type TemplatePrompt struct {
	tmp       *template.Template
	paramName []string
}

func (t *TemplatePrompt) Init(tmp string) error {
	var err error
	t.tmp, err = template.New("template").Parse(tmp)
	if err != nil {
		return err
	}

	// 提取 {{XXX}} 表达式
	re := regexp.MustCompile(`{{(.*?)}}`)
	vRe := regexp.MustCompile(`\.(\w+)`)
	matches := re.FindAllStringSubmatch(tmp, -1)
	paramName := map[string]struct{}{}
	for _, match := range matches {
		if len(match) > 1 {
			// 提取 {{XXX}} 表达式中的变量名
			vMatches := vRe.FindAllStringSubmatch(match[1], -1)
			for _, vMatch := range vMatches {
				paramName[vMatch[1]] = struct{}{}
			}
		}
	}

	t.paramName = make([]string, 0, len(paramName))
	for k := range paramName {
		t.paramName = append(t.paramName, k)
	}
	return err
}

func (t *TemplatePrompt) Format(params utils.ParamsManager) (string, error) {
	if t.tmp == nil {
		return "", fmt.Errorf("模板字段为空")
	}

	// 构造参数
	fParams := make(map[string]string, len(t.paramName))

	for k := range t.paramName {
		v, exist := params.Get(t.paramName[k])
		if exist {
			fParams[t.paramName[k]] = v
		}
	}

	var buff bytes.Buffer
	err := t.tmp.Execute(&buff, fParams)
	if err != nil {
		return "", err
	}
	return buff.String(), nil
}