package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"math"
	"os"
	"sort"
)

type Config struct {
	LabelDimensions map[string]uint64 `json:"label_dimensions"`
	FiniteBuckets   uint64            `json:"finite_buckets"`
	ScrapeSeconds   uint64            `json:"scrape_seconds"`
	SeriesBudget    uint64            `json:"series_budget"`
	RequestsPerSec  float64           `json:"requests_per_second"`
	SpansPerRequest float64           `json:"spans_per_request"`
	TraceKeepRatio  float64           `json:"trace_keep_ratio"`
	BytesPerSpan    float64           `json:"bytes_per_span"`
	LogsPerRequest  float64           `json:"logs_per_request"`
	BytesPerLog     float64           `json:"bytes_per_log"`
	RetentionDays   float64           `json:"retention_days"`
}

type Estimate struct {
	Series, SamplesPerDay           uint64
	TraceRawGiBDay, TraceRawGiBKept float64
	LogRawGiBDay, LogRawGiBKept     float64
}

func decode(r io.Reader) (Config, error) {
	var c Config
	data, err := io.ReadAll(io.LimitReader(r, (64<<10)+1))
	if err != nil {
		return c, err
	}
	if len(data) > 64<<10 {
		return c, errors.New("config exceeds 64 KiB")
	}
	dec := json.NewDecoder(bytes.NewReader(data))
	dec.DisallowUnknownFields()
	if err := dec.Decode(&c); err != nil {
		return c, err
	}
	var extra any
	if err := dec.Decode(&extra); err != io.EOF {
		return c, errors.New("expected exactly one JSON value")
	}
	return c, nil
}

func multiply(a, b uint64) (uint64, error) {
	if b != 0 && a > math.MaxUint64/b {
		return 0, errors.New("integer overflow")
	}
	return a * b, nil
}

func estimate(c Config) (Estimate, error) {
	var result Estimate
	allowed := map[string]bool{"service": true, "instance": true, "method": true, "route": true, "status_class": true}
	if len(c.LabelDimensions) == 0 {
		return result, errors.New("label dimensions must not be empty")
	}
	keys := make([]string, 0, len(c.LabelDimensions))
	for k := range c.LabelDimensions {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	combinations := uint64(1)
	for _, k := range keys {
		if !allowed[k] {
			return result, fmt.Errorf("label %q is not allowed by this model", k)
		}
		if c.LabelDimensions[k] == 0 {
			return result, fmt.Errorf("label %q has zero values", k)
		}
		var err error
		combinations, err = multiply(combinations, c.LabelDimensions[k])
		if err != nil {
			return result, err
		}
	}
	if c.FiniteBuckets == 0 || c.FiniteBuckets > 100 {
		return result, errors.New("finite_buckets must be between 1 and 100")
	}
	if c.ScrapeSeconds == 0 || c.ScrapeSeconds > 86400 || 86400%c.ScrapeSeconds != 0 {
		return result, errors.New("scrape_seconds must divide 86400 exactly")
	}
	series, err := multiply(combinations, c.FiniteBuckets+3) // +Inf, _sum, _count.
	if err != nil {
		return result, err
	}
	if c.SeriesBudget == 0 || series > c.SeriesBudget {
		return result, fmt.Errorf("series budget exceeded: need=%d budget=%d", series, c.SeriesBudget)
	}
	samples, err := multiply(series, 86400/c.ScrapeSeconds)
	if err != nil {
		return result, err
	}
	values := []float64{c.RequestsPerSec, c.SpansPerRequest, c.BytesPerSpan, c.LogsPerRequest, c.BytesPerLog, c.RetentionDays}
	for _, v := range values {
		if math.IsNaN(v) || math.IsInf(v, 0) || v <= 0 {
			return result, errors.New("rates, sizes and retention must be finite and positive")
		}
	}
	if math.IsNaN(c.TraceKeepRatio) || c.TraceKeepRatio < 0 || c.TraceKeepRatio > 1 {
		return result, errors.New("trace_keep_ratio must be in [0,1]")
	}
	trace := c.RequestsPerSec * c.SpansPerRequest * c.TraceKeepRatio * c.BytesPerSpan * 86400 / (1 << 30)
	logs := c.RequestsPerSec * c.LogsPerRequest * c.BytesPerLog * 86400 / (1 << 30)
	result = Estimate{series, samples, trace, trace * c.RetentionDays, logs, logs * c.RetentionDays}
	for _, v := range []float64{result.TraceRawGiBDay, result.TraceRawGiBKept, result.LogRawGiBDay, result.LogRawGiBKept} {
		if math.IsInf(v, 0) || math.IsNaN(v) {
			return Estimate{}, errors.New("floating-point overflow")
		}
	}
	return result, nil
}

func run(path string) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	c, err := decode(f)
	err = errors.Join(err, f.Close())
	if err != nil {
		return err
	}
	e, err := estimate(c)
	if err != nil {
		return err
	}
	fmt.Printf("series=%d samples_per_day=%d\n", e.Series, e.SamplesPerDay)
	fmt.Printf("trace_raw_gib_per_day=%.2f retained_raw_gib=%.2f\n", e.TraceRawGiBDay, e.TraceRawGiBKept)
	fmt.Printf("log_raw_gib_per_day=%.2f retained_raw_gib=%.2f\n", e.LogRawGiBDay, e.LogRawGiBKept)
	fmt.Println("model_only=true network=none backend_verified=false")
	return nil
}

func main() {
	config := flag.String("config", "config.json", "local synthetic sizing assumptions")
	flag.Parse()
	if err := run(*config); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}