package budget

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"time"
)

type watcher struct {
	config config
	ctl    unitCtl
	procs  procSource
	cgroup string
	owned  map[string]bool
	out    io.Writer
}

func randomID() (string, error) {
	var data [3]byte
	_, err := rand.Read(data[:])
	return hex.EncodeToString(data[:]), err
}

func scopeProps(c config, name, slice string, pids []uint32) ([]property, error) {
	if slice == "" {
		slice = "app.slice"
		if name == "search" {
			slice = "search.slice"
		}
	}
	values := []property{prop("Description", name+" workload"), prop("Slice", slice), prop("CollectMode", "inactive-or-failed"), prop("PIDs", pids)}
	if name != "search" {
		limits, err := c.properties(name)
		if err != nil {
			return nil, err
		}
		values = append(values, limits...)
	}
	return values, nil
}

func (w *watcher) init() error {
	units, err := w.ctl.Units()
	if err != nil {
		return err
	}
	w.owned = map[string]bool{}
	for _, unit := range units {
		if strings.HasPrefix(unit.Name, "budget-") && strings.HasSuffix(unit.Name, ".scope") {
			w.owned[unit.Name] = true
		}
	}
	return nil
}

func (w *watcher) reap() error {
	for _, name := range keys(w.owned) {
		state, err := w.ctl.State(name)
		if busError(err, noUnit) {
			delete(w.owned, name)
			continue
		}
		if err != nil {
			return err
		}
		switch state.Active {
		case "inactive", "failed":
			delete(w.owned, name)
		case "active":
			data, err := os.ReadFile(filepath.Join(w.cgroup, strings.TrimPrefix(state.Group, "/"), "cgroup.events"))
			if os.IsNotExist(err) {
				delete(w.owned, name)
				continue
			}
			if err != nil {
				return err
			}
			if strings.Contains(string(data), "populated 0") {
				if err := w.ctl.Stop(name); err != nil && !busError(err, noUnit) {
					return err
				}
				delete(w.owned, name)
			}
		}
	}
	return nil
}

func (w *watcher) step() error {
	tasks, err := w.procs.Snapshot(w.config)
	if err != nil {
		return err
	}
	pids := make([]uint32, 0, len(tasks))
	for pid := range tasks {
		pids = append(pids, pid)
	}
	slices.Sort(pids)
	for _, pid := range pids {
		value := tasks[pid]
		if value.Profile == "" {
			continue
		}
		group, err := w.procs.Group(pid)
		if gone(err) {
			continue
		}
		if err != nil {
			return err
		}
		if group != value.Group {
			continue
		}
		members := family(tasks, pid)
		values, err := scopeProps(w.config, value.Profile, value.Slice, members)
		if err != nil {
			return err
		}
		id, err := randomID()
		if err != nil {
			return err
		}
		name := "budget-" + value.Profile[:1] + id + ".scope"
		if err = w.ctl.Start(name, values); busError(err, noPID) {
			continue
		} else if err != nil {
			return err
		}
		w.owned[name] = true
		if _, err := fmt.Fprintf(w.out, "%s: %s %v\n", value.Profile, name, members); err != nil {
			return err
		}
	}
	return w.reap()
}

func (w *watcher) watch() error {
	if err := w.init(); err != nil {
		return err
	}
	for {
		if err := w.step(); err != nil {
			return err
		}
		time.Sleep(200 * time.Millisecond)
	}
}