package billing

import (
	"context"
	"sync"
	"time"

	"go.uber.org/zap"

	"github.com/mem9-ai/drive9/pkg/logger"
)

// LeaderChecker reports whether this process currently holds the billing
// leader lock. Satisfied by *leader.Manager.
type LeaderChecker interface {
	IsLeader() bool
}

// Worker runs the three billing loops (storage snapshot, hour collection,
// outbox send). All loops are leader-gated: only the leader holder meters
// and sends, so multiple replicas are safe without extra coordination.
type Worker struct {
	cfg       Config
	collector *Collector
	sender    *Sender
	leader    LeaderChecker
}

// NewWorker wires the billing worker. leader may be nil, which means
// always-leader (single-replica deployments).
func NewWorker(cfg Config, store MetaStore, sqsClient SQSSender, owner string, leader LeaderChecker) *Worker {
	cfg = cfg.withDefaults()
	return &Worker{
		cfg:       cfg,
		collector: newCollector(cfg, store, nil),
		sender:    newSender(cfg, store, sqsClient, owner),
		leader:    leader,
	}
}

// Run starts the loops and blocks until ctx is canceled.
func (w *Worker) Run(ctx context.Context) error {
	var wg sync.WaitGroup
	wg.Add(3)
	go func() {
		defer wg.Done()
		w.loop(ctx, "snapshot", w.cfg.SnapshotInterval, func(ctx context.Context) error {
			if err := w.collector.SnapshotOnce(ctx); err != nil {
				return err
			}
			return w.collector.PruneOnce(ctx)
		})
	}()
	go func() {
		defer wg.Done()
		w.loop(ctx, "collect", w.cfg.CollectInterval, func(ctx context.Context) error {
			if err := w.collector.CollectDueHours(ctx); err != nil {
				return err
			}
			return w.collector.RetryIncompleteHours(ctx)
		})
	}()
	go func() {
		defer wg.Done()
		w.loop(ctx, "send", w.cfg.SendInterval, func(ctx context.Context) error {
			defer func() { w.sender.ReportOutboxStats(ctx) }()
			return w.sender.SendOnce(ctx)
		})
	}()
	wg.Wait()
	return ctx.Err()
}

func (w *Worker) loop(ctx context.Context, name string, interval time.Duration, fn func(ctx context.Context) error) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		if ctx.Err() != nil {
			return
		}
		if w.leader == nil || w.leader.IsLeader() {
			if err := fn(ctx); err != nil && ctx.Err() == nil {
				logger.Warn(ctx, "billing_loop_failed",
					zap.String("loop", name), zap.Error(err))
			}
		}
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
	}
}