package osimage

import (
	"k8s.io/klog/v2"
	"sync"
	"time"
)

type progressState struct {
	lastTime    time.Time
	lastPercent int64
	lastBucket  int64
}

type ProgressLogger struct {
	minInterval time.Duration
	bucketSize  int64
	states      map[string]*progressState
}

func NewProgressLogger(minSeconds int, bucketSize int64) *ProgressLogger {
	if minSeconds < 0 {
		minSeconds = 0
	}
	if bucketSize <= 0 {
		bucketSize = 10
	}

	return &ProgressLogger{
		minInterval: time.Duration(minSeconds) * time.Second,
		bucketSize:  bucketSize,
		states:      make(map[string]*progressState),
	}
}

func (l *ProgressLogger) state(stage string) *progressState {
	s, ok := l.states[stage]
	if ok {
		return s
	}
	s = &progressState{
		lastPercent: -1,
		lastBucket:  -1,
	}
	l.states[stage] = s
	return s
}

func (l *ProgressLogger) Log(p Progress) {
	if p.BytesTotal <= 0 {
		return
	}

	percent := PercentOf(p.BytesComplete, p.BytesTotal)

	now := time.Now()
	bucket := percent / l.bucketSize
	s := l.state(p.Stage)

	// Always log first visible progress
	if s.lastPercent == -1 {
		s.lastPercent = percent
		s.lastBucket = bucket
		s.lastTime = now
		klog.V(4).InfoS(p.Stage, "progress", percent)
		return
	}

	// Always log completion once
	if percent == 100 && s.lastPercent < 100 {
		s.lastPercent = 100
		s.lastBucket = 100 / l.bucketSize
		s.lastTime = now
		klog.V(4).InfoS(p.Stage, "progress", 100)
		return
	}

	// Log if we crossed a new milestone bucket
	if bucket > s.lastBucket {
		s.lastPercent = percent
		s.lastBucket = bucket
		s.lastTime = now
		klog.V(4).InfoS(p.Stage, "progress", percent)
		return
	}

	// Otherwise allow a timed refresh if progress moved
	if now.Sub(s.lastTime) >= l.minInterval && percent > s.lastPercent {
		s.lastPercent = percent
		s.lastTime = now
		klog.V(4).InfoS(p.Stage, "progress", percent)
	}
}

type TimeBasedUpdater struct {
	mu       sync.Mutex
	interval time.Duration
	lastRun  time.Time
	inFlight bool
}

func NewTimeBasedUpdater(seconds int) *TimeBasedUpdater {
	if seconds <= 0 {
		seconds = 15
	}
	return &TimeBasedUpdater{
		interval: time.Duration(seconds) * time.Second,
	}
}

func (u *TimeBasedUpdater) Run(fn func() error) error {
	u.mu.Lock()
	now := time.Now()

	if u.inFlight {
		u.mu.Unlock()
		return nil
	}

	if !u.lastRun.IsZero() && now.Sub(u.lastRun) < u.interval {
		u.mu.Unlock()
		return nil
	}

	u.lastRun = now
	u.inFlight = true
	u.mu.Unlock()

	defer func() {
		u.mu.Lock()
		u.inFlight = false
		u.mu.Unlock()
	}()

	return fn()
}