package main

import (
	"context"
	"flag"
	"log/slog"
	"os"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/runner"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

func main() {
	apply := flag.Bool("apply", false, "persist compacted results; default is dry-run")
	batchSize := flag.Int("batch-size", 100, "rows per batch, maximum 100")
	maxBatches := flag.Int("max-batches", 10, "maximum batches for one invocation")
	afterID := flag.String("after-id", "", "resume after this task UUID")
	flag.Parse()

	logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
	cfg := config.Load()
	if err := cfg.Validate(); err != nil {
		logger.Error("invalid gateway configuration", "error", err)
		os.Exit(1)
	}
	if *batchSize < 1 || *batchSize > 100 || *maxBatches < 1 {
		logger.Error("invalid backfill bounds", "batchSize", *batchSize, "maxBatches", *maxBatches)
		os.Exit(1)
	}

	ctx := context.Background()
	db, err := store.Connect(ctx, cfg.DatabaseURL)
	if err != nil {
		logger.Error("connect postgres failed", "error", err)
		os.Exit(1)
	}
	defer db.Close()

	service := runner.New(cfg, db, logger)
	cursor := *afterID
	scanned := 0
	matched := 0
	updated := 0
	expired := 0
	for batch := 0; batch < *maxBatches; batch++ {
		items, err := db.ListTaskBinaryResultBackfillBatch(ctx, cursor, *batchSize)
		if err != nil {
			logger.Error("list binary result backfill batch failed", "afterId", cursor, "error", err)
			os.Exit(1)
		}
		if len(items) == 0 {
			break
		}
		for _, item := range items {
			cursor = item.ID
			scanned++
			if !runner.TaskResultHasInlineBinary(item.Result) {
				continue
			}
			matched++
			if !*apply {
				continue
			}
			isExpired := item.FinishedAt.Before(time.Now().Add(-time.Duration(localResultTTLHours(cfg)) * time.Hour))
			var persistent map[string]any
			var changed bool
			if isExpired {
				persistent, changed, err = service.CompactExpiredTaskResultForStorage(ctx, item.ID, item.Result)
			} else {
				persistent, changed, err = service.MaterializeTaskResultForStorage(ctx, item.ID, item.Result)
			}
			if err != nil {
				logger.Error("materialize historical binary result failed", "taskId", item.ID, "error", err)
				os.Exit(1)
			}
			if !changed {
				continue
			}
			ok, err := db.UpdateTaskBinaryResultBackfill(ctx, item.ID, persistent)
			if err != nil {
				logger.Error("update historical binary result failed", "taskId", item.ID, "error", err)
				os.Exit(1)
			}
			if !ok {
				continue
			}
			updated++
			if isExpired {
				expired++
			}
		}
		if len(items) < *batchSize {
			break
		}
	}
	logger.Info("binary result backfill completed",
		"apply", *apply,
		"scanned", scanned,
		"matched", matched,
		"updated", updated,
		"expired", expired,
		"resumeAfterId", cursor,
	)
}

func localResultTTLHours(cfg config.Config) int {
	if cfg.LocalResultTTLHours <= 0 {
		return 24
	}
	return cfg.LocalResultTTLHours
}