descrybe/apps/api/internal/aiprompts/service.go

package aiprompts

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/descrybe/descrybe-v2/apps/api/internal/company"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// Service loads and stores per-company AI prompt templates.
type Service struct {
	Pool *pgxpool.Pool
}

func NewService(pool *pgxpool.Pool) *Service {
	return &Service{Pool: pool}
}

type stored struct {
	key            string
	language       string
	systemTemplate string
	userTemplate   string
	isEnabled      bool
	updatedAt      time.Time
}

// GetBundle returns all prompt keys with effective templates for language + variable catalog.
func (s *Service) GetBundle(ctx context.Context, companyID uuid.UUID, language string) (Bundle, error) {
	lang, err := company.ParseLanguage(language, true)
	if err != nil {
		lang = company.LoadLanguage(ctx, s.Pool, companyID)
	}
	contentLangs := company.LoadContentLanguages(ctx, s.Pool, companyID)
	storedRows, err := s.loadAll(ctx, companyID)
	if err != nil {
		return Bundle{}, err
	}
	byKeyLang := map[string]stored{}
	customLangs := map[string][]string{}
	for _, st := range storedRows {
		byKeyLang[st.key+"\x00"+st.language] = st
		customLangs[st.key] = appendUnique(customLangs[st.key], st.language)
	}
	out := make([]Template, 0, len(BuiltInDefaults))
	for _, def := range BuiltInDefaults {
		t := Template{
			Key:         def.Key,
			Language:    lang,
			Label:       def.Label,
			Description: def.Description,
			IsDefault:   true,
			IsCustom:    false,
			IsEnabled:   true,
		}
		if st, ok := byKeyLang[def.Key+"\x00"+lang]; ok {
			t.IsCustom = true
			t.IsDefault = false
			t.IsEnabled = st.isEnabled
			t.UpdatedAt = st.updatedAt
			if st.isEnabled {
				t.SystemTemplate = st.systemTemplate
				t.UserTemplate = st.userTemplate
			} else {
				t.SystemTemplate = def.SystemTemplate
				t.UserTemplate = def.UserTemplate
				t.IsDefault = true
			}
		} else {
			t.SystemTemplate = def.SystemTemplate
			t.UserTemplate = def.UserTemplate
		}
		out = append(out, t)
	}
	return Bundle{
		Language:         lang,
		Prompts:          out,
		Variables:        VariableCatalog,
		CustomLanguages:  customLangs,
		ContentLanguages: contentLangs,
	}, nil
}

// Resolve returns the effective templates for one key + language.
// Fallback: requested lang → "*" → company primary language → built-in.
// Language on the result is always the requested (normalized) lang so {{language}}
// still resolves to the content language being generated — templates are shared,
// not duplicated per language.
func (s *Service) Resolve(ctx context.Context, companyID uuid.UUID, key, language string) (Resolved, error) {
	if !ValidPromptKey(key) {
		return Resolved{}, ErrInvalidKey
	}
	def, ok := DefaultFor(key)
	if !ok {
		return Resolved{}, ErrInvalidKey
	}
	lang, err := company.ParseLanguage(language, true)
	if err != nil {
		lang = company.DefaultLanguage
	}
	primary := company.LoadLanguage(ctx, s.Pool, companyID)
	candidates := []string{lang, company.LangPromptAny}
	if primary != "" && primary != lang && primary != company.LangPromptAny {
		candidates = append(candidates, primary)
	}
	for _, cand := range candidates {
		st, err := s.loadOne(ctx, companyID, key, cand)
		if err != nil && !errors.Is(err, pgx.ErrNoRows) {
			return Resolved{}, err
		}
		if err != nil || !st.isEnabled {
			continue
		}
		sys := strings.TrimSpace(st.systemTemplate)
		user := strings.TrimSpace(st.userTemplate)
		if sys == "" {
			sys = def.SystemTemplate
		}
		if user == "" {
			user = def.UserTemplate
		}
		return Resolved{
			Key:            key,
			Language:       lang,
			SystemTemplate: sys,
			UserTemplate:   user,
			IsCustom:       true,
		}, nil
	}
	return Resolved{
		Key:            key,
		Language:       lang,
		SystemTemplate: def.SystemTemplate,
		UserTemplate:   def.UserTemplate,
		IsCustom:       false,
	}, nil
}

// Update applies prompt updates (upsert or reset). Empty prompts list is a no-op.
func (s *Service) Update(ctx context.Context, companyID uuid.UUID, in UpdateInput) (Bundle, error) {
	defaultLang := strings.TrimSpace(in.Language)
	if defaultLang == "" {
		defaultLang = company.LoadLanguage(ctx, s.Pool, companyID)
	}
	if len(in.Prompts) == 0 {
		return s.GetBundle(ctx, companyID, defaultLang)
	}
	tx, err := s.Pool.Begin(ctx)
	if err != nil {
		return Bundle{}, err
	}
	defer tx.Rollback(ctx)
	lastLang := defaultLang
	for _, item := range in.Prompts {
		key := strings.TrimSpace(strings.ToLower(item.Key))
		if !ValidPromptKey(key) {
			return Bundle{}, fmt.Errorf("%w: %s", ErrInvalidKey, item.Key)
		}
		langRaw := strings.TrimSpace(item.Language)
		if langRaw == "" {
			langRaw = defaultLang
		}
		var lang string
		if langRaw == company.LangPromptAny {
			lang = company.LangPromptAny
		} else {
			parsed, err := company.ParseLanguage(langRaw, false)
			if err != nil {
				return Bundle{}, fmt.Errorf("%w: language %q", ErrInvalidInput, langRaw)
			}
			lang = parsed
		}
		lastLang = lang
		if item.Reset {
			_, err := tx.Exec(ctx, `
				DELETE FROM ai_prompt_templates
				WHERE company_id = $1 AND prompt_key = $2 AND language = $3`,
				companyID, key, lang)
			if err != nil {
				return Bundle{}, err
			}
			continue
		}
		sys := SanitizeTemplate(item.SystemTemplate, MaxSystemRunes)
		user := SanitizeTemplate(item.UserTemplate, MaxUserRunes)
		if sys == "" && user == "" {
			return Bundle{}, fmt.Errorf("%w: empty templates for %s", ErrInvalidInput, key)
		}
		enabled := true
		if item.IsEnabled != nil {
			enabled = *item.IsEnabled
		}
		_, err = tx.Exec(ctx, `
			INSERT INTO ai_prompt_templates (
				company_id, prompt_key, language, system_template, user_template, is_enabled, updated_at
			) VALUES ($1,$2,$3,$4,$5,$6, now())
			ON CONFLICT (company_id, prompt_key, language) DO UPDATE SET
				system_template = EXCLUDED.system_template,
				user_template = EXCLUDED.user_template,
				is_enabled = EXCLUDED.is_enabled,
				updated_at = now()`,
			companyID, key, lang, sys, user, enabled)
		if err != nil {
			return Bundle{}, err
		}
	}
	if err := tx.Commit(ctx); err != nil {
		return Bundle{}, err
	}
	bundleLang := lastLang
	if bundleLang == company.LangPromptAny {
		bundleLang = defaultLang
	}
	return s.GetBundle(ctx, companyID, bundleLang)
}

func (s *Service) loadAll(ctx context.Context, companyID uuid.UUID) ([]stored, error) {
	rows, err := s.Pool.Query(ctx, `
		SELECT prompt_key, language, system_template, user_template, is_enabled, updated_at
		FROM ai_prompt_templates WHERE company_id = $1`, companyID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []stored
	for rows.Next() {
		var st stored
		if err := rows.Scan(&st.key, &st.language, &st.systemTemplate, &st.userTemplate, &st.isEnabled, &st.updatedAt); err != nil {
			return nil, err
		}
		out = append(out, st)
	}
	return out, rows.Err()
}

func (s *Service) loadOne(ctx context.Context, companyID uuid.UUID, key, language string) (stored, error) {
	var st stored
	err := s.Pool.QueryRow(ctx, `
		SELECT prompt_key, language, system_template, user_template, is_enabled, updated_at
		FROM ai_prompt_templates
		WHERE company_id = $1 AND prompt_key = $2 AND language = $3`,
		companyID, key, language).Scan(&st.key, &st.language, &st.systemTemplate, &st.userTemplate, &st.isEnabled, &st.updatedAt)
	return st, err
}

func appendUnique(list []string, v string) []string {
	for _, x := range list {
		if x == v {
			return list
		}
	}
	return append(list, v)
}