package internal

import (
	"context"
	"log"
	"runtime"
	"sync/atomic"
	"time"
)

const MonitorInterval = 30 * time.Second

type MemoryMonitor struct {
	peakAlloc atomic.Uint64
}

func NewMemoryMonitor() *MemoryMonitor {
	return &MemoryMonitor{}
}

// Snapshot reads current heap allocation, updates the peak, and returns both.
func (m *MemoryMonitor) Snapshot() (currentAlloc, peakAlloc uint64) {
	var ms runtime.MemStats
	runtime.ReadMemStats(&ms)
	currentAlloc = ms.HeapAlloc
	for {
		peak := m.peakAlloc.Load()
		if currentAlloc <= peak {
			return currentAlloc, peak
		}
		if m.peakAlloc.CompareAndSwap(peak, currentAlloc) {
			return currentAlloc, currentAlloc
		}
	}
}

// Run starts a periodic monitor loop that calls logFn with each snapshot.
// It blocks until ctx is cancelled.
func (m *MemoryMonitor) Run(ctx context.Context, logFn func(currentAlloc, peakAlloc uint64)) {
	ticker := time.NewTicker(MonitorInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			current, peak := m.Snapshot()
			logFn(current, peak)
		}
	}
}

// DefaultLogFn returns a log function with the given prefix.
func DefaultLogFn(prefix string) func(currentAlloc, peakAlloc uint64) {
	return func(currentAlloc, peakAlloc uint64) {
		log.Printf("[%s] heap alloc: %s | peak heap alloc: %s",
			prefix, FormatBytes(currentAlloc), FormatBytes(peakAlloc))
	}
}