package execution

import (
	"context"
	"errors"
	"log/slog"
	"sort"
	"strings"
	"time"

	executionstore "github.com/tidbcloud/AgentTeam/internal/store/role/execution"
)

const (
	defaultWorkerPollInterval = 5 * time.Second
	defaultWorkerLeaseTTL     = 2 * time.Minute
	defaultWorkerClaimLimit   = 1
)

type WorkerOptions struct {
	WorkspaceIDs  func(context.Context) ([]string, error)
	ClaimRunnable func(
		context.Context,
		executionstore.ClaimParams,
	) ([]executionstore.ClaimedExecution, error)
	Execute func(
		context.Context,
		executionstore.ClaimedExecution,
		time.Time,
	) error
	Logger   *slog.Logger
	Now      func() time.Time
	Interval time.Duration
	LeaseTTL time.Duration
	Limit    int
}

type Worker struct {
	workspaceIDs  func(context.Context) ([]string, error)
	claimRunnable func(
		context.Context,
		executionstore.ClaimParams,
	) ([]executionstore.ClaimedExecution, error)
	execute func(
		context.Context,
		executionstore.ClaimedExecution,
		time.Time,
	) error
	logger   *slog.Logger
	now      func() time.Time
	interval time.Duration
	leaseTTL time.Duration
	limit    int
}

func NewWorker(options WorkerOptions) (*Worker, error) {
	if options.WorkspaceIDs == nil ||
		options.ClaimRunnable == nil ||
		options.Execute == nil {
		return nil, errors.New("role management worker dependencies are incomplete")
	}
	if options.Now == nil {
		options.Now = time.Now
	}
	if options.Interval <= 0 {
		options.Interval = defaultWorkerPollInterval
	}
	if options.LeaseTTL <= 0 {
		options.LeaseTTL = defaultWorkerLeaseTTL
	}
	if options.Limit <= 0 {
		options.Limit = defaultWorkerClaimLimit
	}
	return &Worker{
		workspaceIDs: options.WorkspaceIDs, claimRunnable: options.ClaimRunnable,
		execute: options.Execute, logger: options.Logger, now: options.Now,
		interval: options.Interval, leaseTTL: options.LeaseTTL, limit: options.Limit,
	}, nil
}

func (worker *Worker) Run(ctx context.Context) {
	if worker == nil || ctx.Err() != nil {
		return
	}
	worker.ProcessOnce(ctx)
	ticker := time.NewTicker(worker.interval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			worker.ProcessOnce(ctx)
		}
	}
}

func (worker *Worker) ProcessOnce(ctx context.Context) {
	if worker == nil || ctx.Err() != nil {
		return
	}
	workspaceIDs, err := worker.workspaceIDs(ctx)
	if err != nil {
		worker.warn(ctx, "list workspaces", err)
		return
	}
	if ctx.Err() != nil {
		return
	}
	for _, workspaceID := range canonicalWorkspaceIDs(workspaceIDs) {
		if ctx.Err() != nil {
			return
		}
		claimNow := worker.now().UTC()
		claimed, err := worker.claimRunnable(ctx, executionstore.ClaimParams{
			WorkspaceID: workspaceID,
			Now:         claimNow,
			LeaseTTL:    worker.leaseTTL,
			Limit:       worker.limit,
		})
		if ctx.Err() != nil {
			return
		}
		if err != nil {
			worker.warn(ctx, "claim", err)
			continue
		}
		for _, claim := range claimed {
			if ctx.Err() != nil {
				return
			}
			err := worker.execute(ctx, claim, claimNow)
			if ctx.Err() != nil {
				return
			}
			if err != nil {
				worker.warn(ctx, "execute", err)
			}
		}
	}
}

func canonicalWorkspaceIDs(values []string) []string {
	unique := make(map[string]struct{}, len(values))
	for _, value := range values {
		value = strings.TrimSpace(value)
		if value != "" {
			unique[value] = struct{}{}
		}
	}
	result := make([]string, 0, len(unique))
	for value := range unique {
		result = append(result, value)
	}
	sort.Strings(result)
	return result
}

func (worker *Worker) warn(ctx context.Context, operation string, err error) {
	if worker.logger != nil {
		worker.logger.WarnContext(
			ctx, "role management worker: "+operation, "error", err,
		)
	}
}