geo/server/internal/tenant/app/schedule_task_policy.go

package app

import (
	"encoding/json"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/geo-platform/tenant-api/internal/shared/response"
	sharedschedule "github.com/geo-platform/tenant-api/internal/shared/schedule"
)

const (
	ScheduleTargetPromptRule            = "prompt_rule"
	ScheduleTargetKolSubscriptionPrompt = "kol_subscription_prompt"

	scheduleKindDaily  = "daily"
	scheduleKindWeekly = "weekly"

	defaultRandomWindowStart = "01:00:00"
	defaultRandomWindowEnd   = "05:00:00"
)

var scheduleDayOrder = []string{"mon", "tue", "wed", "thu", "fri", "sat", "sun"}

type normalizedSchedulePlan struct {
	CronExpr          string
	ScheduleKind      string
	ScheduleDays      []string
	ScheduleTimeMode  string
	RandomWindowStart string
	RandomWindowEnd   string
	windowStart       time.Duration
	windowEnd         time.Duration
}

func normalizeScheduleTarget(req ScheduleTaskRequest) (string, *int64, *int64, error) {
	targetType := strings.TrimSpace(req.TargetType)
	if targetType == "" {
		switch {
		case req.SubscriptionPromptID != nil && *req.SubscriptionPromptID > 0:
			targetType = ScheduleTargetKolSubscriptionPrompt
		default:
			targetType = ScheduleTargetPromptRule
		}
	}

	switch targetType {
	case ScheduleTargetPromptRule:
		if req.PromptRuleID == nil || *req.PromptRuleID <= 0 {
			return "", nil, nil, response.ErrBadRequest(40001, "invalid_rule", "prompt_rule_id is required")
		}
		promptRuleID := *req.PromptRuleID
		return targetType, &promptRuleID, nil, nil
	case ScheduleTargetKolSubscriptionPrompt:
		if req.SubscriptionPromptID == nil || *req.SubscriptionPromptID <= 0 {
			return "", nil, nil, response.ErrBadRequest(40075, "invalid_subscription_prompt", "subscription_prompt_id is required")
		}
		subscriptionPromptID := *req.SubscriptionPromptID
		return targetType, nil, &subscriptionPromptID, nil
	default:
		return "", nil, nil, response.ErrBadRequest(40076, "invalid_schedule_target", "target_type must be prompt_rule or kol_subscription_prompt")
	}
}

func normalizeSchedulePlan(req ScheduleTaskRequest) (normalizedSchedulePlan, error) {
	kind := strings.TrimSpace(req.ScheduleKind)
	if kind == "" {
		kind = scheduleKindDaily
	}
	if kind != scheduleKindDaily && kind != scheduleKindWeekly {
		return normalizedSchedulePlan{}, response.ErrBadRequest(40077, "invalid_schedule_kind", "schedule_kind must be daily or weekly")
	}

	dayKeys, err := normalizeScheduleDayKeys(kind, req.ScheduleDays)
	if err != nil {
		return normalizedSchedulePlan{}, err
	}

	timeMode := strings.TrimSpace(req.ScheduleTimeMode)
	if timeMode == "" {
		if strings.TrimSpace(req.CronExpr) != "" {
			timeMode = sharedschedule.TimeModeFixed
		} else {
			timeMode = sharedschedule.TimeModeRandomWindow
		}
	}
	if timeMode != sharedschedule.TimeModeFixed && timeMode != sharedschedule.TimeModeRandomWindow {
		return normalizedSchedulePlan{}, response.ErrBadRequest(40078, "invalid_schedule_time_mode", "schedule_time_mode must be fixed or random_window")
	}

	windowStart := strings.TrimSpace(req.RandomWindowStart)
	if windowStart == "" {
		windowStart = defaultRandomWindowStart
	}
	windowEnd := strings.TrimSpace(req.RandomWindowEnd)
	if windowEnd == "" {
		windowEnd = defaultRandomWindowEnd
	}
	windowStartDuration, err := sharedschedule.ParseClockDuration(windowStart)
	if err != nil {
		return normalizedSchedulePlan{}, response.ErrBadRequest(40079, "invalid_random_window_start", err.Error())
	}
	windowEndDuration, err := sharedschedule.ParseClockDuration(windowEnd)
	if err != nil {
		return normalizedSchedulePlan{}, response.ErrBadRequest(40080, "invalid_random_window_end", err.Error())
	}
	if windowEndDuration <= windowStartDuration {
		return normalizedSchedulePlan{}, response.ErrBadRequest(40081, "invalid_random_window", "random_window_end must be after random_window_start")
	}

	cronExpr := strings.TrimSpace(req.CronExpr)
	if timeMode == sharedschedule.TimeModeFixed {
		if cronExpr == "" {
			cronExpr = buildScheduleCronFromDays(dayKeys, windowStartDuration)
		}
		if err := validateScheduleCronExpr(cronExpr); err != nil {
			return normalizedSchedulePlan{}, err
		}
	} else if cronExpr == "" {
		cronExpr = buildScheduleCronFromDays(dayKeys, windowStartDuration)
	}

	return normalizedSchedulePlan{
		CronExpr:          cronExpr,
		ScheduleKind:      kind,
		ScheduleDays:      dayKeys,
		ScheduleTimeMode:  timeMode,
		RandomWindowStart: sharedschedule.FormatClockDuration(windowStartDuration),
		RandomWindowEnd:   sharedschedule.FormatClockDuration(windowEndDuration),
		windowStart:       windowStartDuration,
		windowEnd:         windowEndDuration,
	}, nil
}

func normalizeScheduleDayKeys(kind string, values []string) ([]string, error) {
	if kind == scheduleKindDaily {
		return append([]string(nil), scheduleDayOrder...), nil
	}
	if len(values) == 0 {
		return nil, response.ErrBadRequest(40082, "schedule_days_required", "schedule_days is required for weekly schedule")
	}

	seen := make(map[string]struct{}, len(values))
	for _, value := range values {
		day, ok := sharedschedule.ParseDayKey(value)
		if !ok {
			return nil, response.ErrBadRequest(40083, "invalid_schedule_day", fmt.Sprintf("invalid schedule day: %s", value))
		}
		key := sharedschedule.DayKey(day)
		if key == "" {
			return nil, response.ErrBadRequest(40083, "invalid_schedule_day", fmt.Sprintf("invalid schedule day: %s", value))
		}
		seen[key] = struct{}{}
	}
	if len(seen) == 0 {
		return nil, response.ErrBadRequest(40082, "schedule_days_required", "schedule_days is required")
	}
	if len(seen) > 7 {
		return nil, response.ErrBadRequest(40084, "too_many_schedule_days", "schedule_days can contain at most 7 days")
	}

	result := make([]string, 0, len(seen))
	for _, key := range scheduleDayOrder {
		if _, ok := seen[key]; ok {
			result = append(result, key)
		}
	}
	return result, nil
}

func buildScheduleCronFromDays(dayKeys []string, at time.Duration) string {
	hour := int(at / time.Hour)
	minute := int((at - time.Duration(hour)*time.Hour) / time.Minute)
	if len(dayKeys) == 0 || len(dayKeys) == 7 {
		return fmt.Sprintf("%d %d * * *", minute, hour)
	}
	return fmt.Sprintf("%d %d * * %s", minute, hour, strings.Join(dayKeys, ","))
}

func scheduleDayKeysToWeekdays(dayKeys []string) []time.Weekday {
	result := make([]time.Weekday, 0, len(dayKeys))
	for _, key := range dayKeys {
		day, ok := sharedschedule.ParseDayKey(key)
		if ok {
			result = append(result, day)
		}
	}
	return result
}

func decodeScheduleDayKeys(raw []byte) []string {
	if len(raw) == 0 {
		return append([]string(nil), scheduleDayOrder...)
	}
	var values []string
	if err := json.Unmarshal(raw, &values); err != nil {
		return append([]string(nil), scheduleDayOrder...)
	}
	kind := scheduleKindWeekly
	if len(values) == 7 {
		kind = scheduleKindDaily
	}
	normalized, err := normalizeScheduleDayKeys(kind, values)
	if err != nil || len(normalized) == 0 {
		return append([]string(nil), scheduleDayOrder...)
	}
	return normalized
}

func DecodeScheduleDayKeys(raw []byte) []string {
	return decodeScheduleDayKeys(raw)
}

func encodeScheduleDayKeys(dayKeys []string) []byte {
	if len(dayKeys) == 0 {
		dayKeys = scheduleDayOrder
	}
	raw, err := json.Marshal(dayKeys)
	if err != nil {
		return []byte(`["mon","tue","wed","thu","fri","sat","sun"]`)
	}
	return raw
}

func decodeScheduleGenerationInput(raw []byte) map[string]interface{} {
	if len(raw) == 0 {
		return map[string]interface{}{}
	}
	var value map[string]interface{}
	if err := json.Unmarshal(raw, &value); err != nil || value == nil {
		return map[string]interface{}{}
	}
	return value
}

func encodeScheduleGenerationInput(value map[string]interface{}) []byte {
	if len(value) == 0 {
		return []byte(`{}`)
	}
	raw, err := json.Marshal(value)
	if err != nil {
		return []byte(`{}`)
	}
	return raw
}

func scheduleGenerationInputCacheKey(value map[string]interface{}) string {
	if len(value) == 0 {
		return "{}"
	}
	keys := make([]string, 0, len(value))
	for key := range value {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	ordered := make(map[string]interface{}, len(keys))
	for _, key := range keys {
		ordered[key] = value[key]
	}
	return digestCacheKey(ordered)
}

func computeScheduleNextRun(
	status string,
	taskID int64,
	cronExpr string,
	dayKeys []string,
	timeMode string,
	windowStart string,
	windowEnd string,
	startAt *time.Time,
	endAt *time.Time,
	now time.Time,
) (*time.Time, error) {
	if status != "enabled" {
		return nil, nil
	}
	if strings.TrimSpace(timeMode) != sharedschedule.TimeModeRandomWindow {
		return sharedschedule.ComputeNextRun(cronExpr, startAt, endAt, now)
	}

	windowStartDuration, err := sharedschedule.ParseClockDuration(windowStart)
	if err != nil {
		return nil, err
	}
	windowEndDuration, err := sharedschedule.ParseClockDuration(windowEnd)
	if err != nil {
		return nil, err
	}
	return sharedschedule.ComputeNextRandomRun(sharedschedule.RandomPolicy{
		Days:        scheduleDayKeysToWeekdays(dayKeys),
		TimeMode:    sharedschedule.TimeModeRandomWindow,
		WindowStart: windowStartDuration,
		WindowEnd:   windowEndDuration,
		Seed:        taskID,
	}, startAt, endAt, now)
}

func ComputeScheduleTaskNextRun(
	status string,
	taskID int64,
	cronExpr string,
	dayKeys []string,
	timeMode string,
	windowStart string,
	windowEnd string,
	startAt *time.Time,
	endAt *time.Time,
	now time.Time,
) (*time.Time, error) {
	return computeScheduleNextRun(status, taskID, cronExpr, dayKeys, timeMode, windowStart, windowEnd, startAt, endAt, now)
}

func normalizeKolScheduleGenerationInput(input map[string]interface{}) map[string]interface{} {
	if input == nil {
		return map[string]interface{}{}
	}
	result := clonePromptRuleExtraParams(input)
	if variables, ok := result["variables"].(map[string]interface{}); ok && variables != nil {
		result["variables"] = variables
	} else {
		result["variables"] = map[string]interface{}{}
	}
	return result
}