package fs2

import (
	"bufio"
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"

	"linkfog.com/public/lib/cgroup/types"
	"linkfog.com/public/lib/l"
)

func SetIORWLimit(cgroupDir string, devinfos []types.DevInfo, rlimit, wlimit uint64) (map[string]*types.DevLimit, error) {
	ioLimitPath := path.Join(cgroupDir, "io.max")
	l.Infof("io limit path: %s, rlimit size: %d, wlimit size: %d", ioLimitPath, rlimit, wlimit)
	f, err := os.OpenFile(ioLimitPath, os.O_WRONLY|os.O_CREATE, 0)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	devLmts := make(map[string]*types.DevLimit, 0)
	ioRdLimitSuccCnt := 0
	ioWrLimitSuccCnt := 0
	for _, devinfo := range devinfos {
		devKey := fmt.Sprintf("%s:%s", devinfo.Major, devinfo.Minor)
		devLmts[devKey] = &types.DevLimit{
			DevName: devinfo.Name,
			RdLimit: false,
			WrLimit: false,
		}

		var ioRdLimit string
		var ioWrLimit string
		if rlimit == 0 { // delete limit
			ioRdLimit = fmt.Sprintf("%s:%s %s=max", devinfo.Major, devinfo.Minor, "rbps")
		} else {
			ioRdLimit = fmt.Sprintf("%s:%s %s=%d", devinfo.Major, devinfo.Minor, "rbps", rlimit)
		}
		if wlimit == 0 { // delete limit
			ioWrLimit = fmt.Sprintf("%s:%s %s=max", devinfo.Major, devinfo.Minor, "wbps")
		} else {
			ioWrLimit = fmt.Sprintf("%s:%s %s=%d", devinfo.Major, devinfo.Minor, "wbps", wlimit)
		}

		// set io read limit
		_, err := f.WriteString(ioRdLimit)
		if err != nil {
			l.Debugf("set %s io read limit error %v, config:%s", devinfo.Name, err, ioRdLimit)
		} else {
			devLmts[devKey].RdLimit = true
			ioRdLimitSuccCnt++
			l.Debugf("set %s io read limit success, config:%s", devinfo.Name, ioRdLimit)
		}

		// set io write limit
		_, err = f.WriteString(ioWrLimit)
		if err != nil {
			l.Debugf("set %s io write limit error %v, config:%s", devinfo.Name, err, ioWrLimit)
		} else {
			devLmts[devKey].WrLimit = true
			ioWrLimitSuccCnt++
			l.Debugf("set %s io write limit success, config:%s", devinfo.Name, ioWrLimit)
		}
	}
	if ioRdLimitSuccCnt == 0 || ioWrLimitSuccCnt == 0 {
		return nil, fmt.Errorf("set io limit error, read_limit:%d, write_limit:%d",
			ioRdLimitSuccCnt, ioWrLimitSuccCnt)
	}

	return devLmts, nil
}

// # cat io.stat
// 253:0 rbytes=0 wbytes=12288 rios=0 wios=3 dbytes=0 dios=0
// 8:0 rbytes=0 wbytes=12288 rios=0 wios=3 dbytes=0 dios=0
func GetAllDevIOStat(cgroupDir string) (map[string]*types.IOStat, error) {
	values, err := parseIOStatFile(filepath.Join(cgroupDir, "io.stat"))
	if err != nil {
		return nil, err
	}

	ioStats := make(map[string]*types.IOStat, 0)
	ioStats["total"] = &types.IOStat{}
	for k, v := range values {
		d := strings.Split(k, ":")
		if len(d) != 2 {
			continue
		}
		device := k
		for _, item := range v {
			d := strings.Split(item, "=")
			if len(d) != 2 {
				continue
			}
			op := d[0]
			if op == "rbytes" || op == "wbytes" {
				sizeUint64, err := strconv.ParseUint(d[1], 10, 64)
				if err != nil {
					return nil, err
				}
				if _, ok := ioStats[device]; !ok {
					ioStats[device] = &types.IOStat{}
				}
				if op == "rbytes" {
					ioStats[device].Read = sizeUint64
					ioStats["total"].Read += sizeUint64
				}
				if op == "wbytes" {
					ioStats[device].Write = sizeUint64
					ioStats["total"].Write += sizeUint64
				}
			}
		}
	}
	return ioStats, nil
}

func parseIOStatFile(ioStatFile string) (map[string][]string, error) {
	f, err := os.OpenFile(ioStatFile, os.O_RDONLY, 0)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	ret := map[string][]string{}
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		parts := strings.Fields(line)
		if len(parts) < 2 {
			continue
		}
		ret[parts[0]] = parts[1:]
	}
	if err := scanner.Err(); err != nil {
		return nil, err
	}
	return ret, nil
}