zarinpal-analytics-backend/cmd/worker/main.go

package main

import (
	"context"
	"errors"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"time"

	"git.amiirkhl.ir/starcoach-hackathon/zarinpal-analytics-backend/internal/aiprovider"
	"git.amiirkhl.ir/starcoach-hackathon/zarinpal-analytics-backend/internal/coach"
	"git.amiirkhl.ir/starcoach-hackathon/zarinpal-analytics-backend/internal/config"
	"git.amiirkhl.ir/starcoach-hackathon/zarinpal-analytics-backend/internal/database"
	"git.amiirkhl.ir/starcoach-hackathon/zarinpal-analytics-backend/internal/growthreport"
	"git.amiirkhl.ir/starcoach-hackathon/zarinpal-analytics-backend/internal/insight"
	"github.com/redis/go-redis/v9"
)

func main() {
	logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
	cfg, err := config.Load()
	if err != nil {
		logger.Error("configuration failed", "error", err)
		os.Exit(1)
	}
	if err := cfg.ValidateAPI(); err != nil {
		logger.Error("worker configuration failed", "error", err)
		os.Exit(1)
	}
	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
	defer cancel()
	pool, err := database.Open(ctx, cfg.DatabaseURL)
	if err != nil {
		logger.Error("database connection failed", "error", err)
		os.Exit(1)
	}
	defer pool.Close()
	redisClient := redis.NewClient(&redis.Options{Addr: cfg.RedisAddr, Password: cfg.RedisPassword})
	defer redisClient.Close()
	repository := growthreport.NewRepository(pool)
	coordinator := growthreport.NewCoordinator(redisClient, growthreport.CoordinatorConfig{CreatePerMinute: cfg.AICreatePerMinute, UncachedPerDay: cfg.AIUncachedScopesPerDay, DailyRequests: cfg.AIDailyRequestBudget, DailyTokens: cfg.AIDailyTokenBudget, MaxConcurrent: cfg.AIMaxConcurrent})
	provider := cfg.SelectedAIProvider()
	var growthGenerator growthreport.CopyGenerator
	var sharedProvider aiprovider.Generator
	if provider.Name == config.AIProviderAgentRouter {
		growthGenerator = growthreport.NewAgentRouterClient(growthreport.AgentRouterConfig{BaseURL: provider.BaseURL, APIKey: provider.APIKey, Model: provider.Model, Timeout: 90 * time.Second}, &http.Client{})
		sharedProvider = aiprovider.NewAgentRouter(aiprovider.Config{BaseURL: provider.BaseURL, APIKey: provider.APIKey, Model: provider.Model, Timeout: 90 * time.Second}, &http.Client{})
	} else {
		growthGenerator = growthreport.NewAvalAIClient(growthreport.AvalAIConfig{BaseURL: provider.BaseURL, APIKey: provider.APIKey, Model: provider.Model, Timeout: 30 * time.Second}, &http.Client{})
		sharedProvider = aiprovider.New(aiprovider.Config{BaseURL: provider.BaseURL, APIKey: provider.APIKey, Model: provider.Model, Timeout: 30 * time.Second}, &http.Client{})
	}
	aiAvailable := provider.APIKey != "" && provider.Model != ""
	processor := growthreport.NewProcessor(repository, growthreport.NewPostgreSQLBuilder(pool), growthGenerator, coordinator, cfg.AIGrowthEnabled && aiAvailable)
	var insightService *insight.Service
	if cfg.AIPageInsightsEnabled && aiAvailable {
		insightService = insight.NewService(insight.NewRepository(pool), insight.NewPostgreSQLBuilder(pool), insight.NewAvalAIGenerator(sharedProvider)).WithBudget(cfg.AIPageInsightsDailyRequestBudget, cfg.AIPageInsightsDailyTokenBudget)
	} else {
		insightService = insight.NewService(insight.NewRepository(pool), insight.NewPostgreSQLBuilder(pool))
	}
	var coachGenerator coach.Generator
	if cfg.AICoachEnabled && aiAvailable {
		coachGenerator = coach.NewAvalAIGenerator(sharedProvider)
	}
	coachService := coach.NewService(coach.NewRepository(pool), coachGenerator).WithBudget(cfg.AICoachDailyRequestBudget, cfg.AICoachDailyTokenBudget)
	logger.Info("growth worker ready", "ai_enabled", cfg.AIGrowthEnabled && aiAvailable, "ai_provider", provider.Name, "ai_model", provider.Model)
	poll := time.NewTicker(time.Second)
	defer poll.Stop()
	cleanup := time.NewTicker(24 * time.Hour)
	defer cleanup.Stop()
	evaluation := time.NewTicker(time.Minute)
	defer evaluation.Stop()
	if cfg.AIPageInsightsEnabled {
		go runInsightScheduler(ctx, logger, insightService)
	}
	for {
		select {
		case <-ctx.Done():
			return
		case <-poll.C:
			var err error
			if cfg.AIPageInsightsEnabled {
				insightCtx, insightCancel := context.WithTimeout(ctx, 2*time.Minute)
				err = insightService.ProcessNext(insightCtx)
				insightCancel()
				if err != nil && !errors.Is(err, insight.ErrNotFound) {
					logger.Error("page insight processing failed", "error", err)
				}
			}
			workCtx, workCancel := context.WithTimeout(ctx, 2*time.Minute)
			err = processor.ProcessNext(workCtx)
			workCancel()
			if err != nil && !errors.Is(err, growthreport.ErrNotFound) {
				logger.Error("growth report processing failed", "error", err)
			}
			if cfg.AICoachEnabled {
				coachCtx, coachCancel := context.WithTimeout(ctx, 2*time.Minute)
				err = coachService.ProcessNext(coachCtx)
				coachCancel()
				if err != nil && !errors.Is(err, coach.ErrNotFound) {
					logger.Error("growth coach processing failed", "error", err)
				}
			}
		case <-cleanup.C:
			cleanupCtx, cleanupCancel := context.WithTimeout(ctx, 30*time.Second)
			deleted, err := repository.CleanupSuperseded(cleanupCtx, time.Now().UTC().AddDate(0, 0, -90))
			cleanupCancel()
			if err != nil {
				logger.Error("growth cleanup failed", "error", err)
			} else if deleted > 0 {
				logger.Info("growth cleanup complete", "deleted_reports", deleted)
			}
		case <-evaluation.C:
			if cfg.AICoachEnabled {
				evaluationCtx, evaluationCancel := context.WithTimeout(ctx, time.Minute)
				completed, err := coachService.EvaluateMetrics(evaluationCtx)
				evaluationCancel()
				if err != nil {
					logger.Error("coach metric evaluation failed", "error", err)
				} else if completed > 0 {
					logger.Info("coach metric goals completed", "count", completed)
				}
				difficultyCtx, difficultyCancel := context.WithTimeout(ctx, time.Minute)
				proposals, err := coachService.EvaluateDifficulty(difficultyCtx)
				difficultyCancel()
				if err != nil {
					logger.Error("coach difficulty evaluation failed", "error", err)
				} else if proposals > 0 {
					logger.Info("coach difficulty proposals created", "count", proposals)
				}
			}
		}
	}
}

func runInsightScheduler(ctx context.Context, logger *slog.Logger, service *insight.Service) {
	run := func() {
		scheduleCtx, cancel := context.WithTimeout(ctx, 30*time.Minute)
		defer cancel()
		created, err := service.EnqueueAll(scheduleCtx, "NIGHTLY")
		if err != nil {
			if !errors.Is(err, context.Canceled) {
				logger.Error("insight scheduling failed", "error", err)
			}
			return
		}
		logger.Info("insight scheduling complete", "created_jobs", created)
	}
	run()
	ticker := time.NewTicker(time.Hour)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			run()
		}
	}
}