package limit

import (
	"io/ioutil"
	"strconv"
	"strings"

	"linkfog.com/public/lib/cgroup"
	"linkfog.com/public/lib/cgroup/fs"
	"linkfog.com/public/lib/cgroup/fs2"
	"linkfog.com/public/lib/cgroup/types"
)

func SetSelfIORWLimit(rlimit, wlimit uint64) (map[string]*types.DevLimit, error) {
	return SetIORWLimit(-1, rlimit, wlimit)
}

func SetIORWLimit(pid int, rlimit, wlimit uint64) (map[string]*types.DevLimit, error) {
	cgroupDir, err := cgroup.GetContainerIOCgroupPath(pid)
	if err != nil {
		return nil, err
	}

	devinfos, err := getDevInfos("/proc/diskstats")
	if err != nil {
		return nil, err
	}

	if cgroup.IsCgroup2() {
		return fs2.SetIORWLimit(cgroupDir, devinfos, rlimit, wlimit)
	} else {
		return fs.SetIORWLimit(cgroupDir, devinfos, rlimit, wlimit)
	}
}

func getDevInfos(file string) ([]types.DevInfo, error) {
	devinfos := make([]types.DevInfo, 0)
	data, err := ioutil.ReadFile(file)
	if err != nil {
		return devinfos, err
	}

	for _, line := range strings.Split(string(data), "\n") {
		line = strings.TrimSpace(line)
		fields := strings.Fields(line)
		if len(fields) >= 3 {
			_, err := strconv.Atoi(fields[0])
			if err != nil {
				continue
			}
			_, err = strconv.Atoi(fields[1])
			if err != nil {
				continue
			}
			devinfos = append(devinfos, types.DevInfo{Major: fields[0], Minor: fields[1], Name: fields[2]})
		}
	}

	return devinfos, nil
}