package main

import (
	"flag"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/collectors"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var (
	allocatedBytes = prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "memconsumer_allocated_bytes",
		Help: "Total bytes allocated by the memory consumer",
	})
	allocationCount = prometheus.NewCounter(prometheus.CounterOpts{
		Name: "memconsumer_allocations_total",
		Help: "Total number of memory allocations performed",
	})
)

// memoryHolder prevents GC from reclaiming allocations
var memoryHolder [][]byte

func envInt(key string, fallback int) int {
	if v := os.Getenv(key); v != "" {
		if i, err := strconv.Atoi(v); err == nil {
			return i
		}
	}
	return fallback
}

func envDuration(key string, fallback time.Duration) time.Duration {
	if v := os.Getenv(key); v != "" {
		if d, err := time.ParseDuration(v); err == nil {
			return d
		}
	}
	return fallback
}

func envString(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}

func main() {
	// Env vars set defaults, flags can override
	initialMB := flag.Int("initial-mb", envInt("INITIAL_MB", 0), "Initial memory to allocate (MB) [env: INITIAL_MB]")
	allocMB := flag.Int("alloc-mb", envInt("ALLOC_MB", 10), "Memory to allocate per iteration (MB) [env: ALLOC_MB]")
	maxMB := flag.Int("max-mb", envInt("MAX_MB", 1000), "Max total memory to allocate, 0=unlimited (MB) [env: MAX_MB]")
	interval := flag.Duration("interval", envDuration("INTERVAL", 5*time.Second), "Allocation interval [env: INTERVAL]")
	metricsAddr := flag.String("metrics-addr", envString("METRICS_ADDR", ":8080"), "Prometheus metrics address [env: METRICS_ADDR]")
	flag.Parse()

	var totalAllocated int

	reg := prometheus.NewRegistry()
	reg.MustRegister(
		collectors.NewGoCollector(),
		collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}),
		allocatedBytes,
		allocationCount,
	)

	// Initial allocation
	if *initialMB > 0 {
		allocate(*initialMB)
		totalAllocated += *initialMB
		log.Printf("Initial allocation: %d MB", *initialMB)
	}

	// Metrics server
	http.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
	http.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	})
	go func() {
		log.Printf("Metrics server listening on %s", *metricsAddr)
		if err := http.ListenAndServe(*metricsAddr, nil); err != nil {
			log.Fatalf("metrics server failed: %v", err)
		}
	}()

	// Graceful shutdown
	stop := make(chan os.Signal, 1)
	signal.Notify(stop, syscall.SIGTERM, syscall.SIGINT)

	ticker := time.NewTicker(*interval)
	defer ticker.Stop()

	log.Printf("Starting: alloc=%dMB interval=%v max=%dMB", *allocMB, *interval, *maxMB)

	for {
		select {
		case <-ticker.C:
			if *maxMB > 0 && totalAllocated+*allocMB > *maxMB {
				log.Printf("Max limit reached (%d MB), skipping allocation", *maxMB)
				continue
			}
			allocate(*allocMB)
			totalAllocated += *allocMB
			log.Printf("Allocated %d MB, total held: %d MB", *allocMB, totalAllocated)
		case <-stop:
			log.Println("Shutting down")
			return
		}
	}
}

func allocate(mb int) {
	size := mb * 1024 * 1024
	buf := make([]byte, size)
	// Touch pages to ensure RSS allocation
	for i := 0; i < size; i += 4096 {
		buf[i] = 1
	}
	memoryHolder = append(memoryHolder, buf)
	allocatedBytes.Add(float64(size))
	allocationCount.Inc()
}