package fs

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

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

func IOSubSystemName() string {
	return "blkio"
}

func SetIORWLimit(cgroupDir string, devinfos []types.DevInfo, rlimit, wlimit uint64) (map[string]*types.DevLimit, error) {
	ioRdLimitPath := path.Join(cgroupDir, "blkio.throttle.read_bps_device")
	ioWrLimitPath := path.Join(cgroupDir, "blkio.throttle.write_bps_device")
	l.Infof("io read limit path: %s, limit size: %d", ioRdLimitPath, rlimit)
	l.Infof("io write limit path: %s, limit size: %d", ioWrLimitPath, wlimit)
	fr, err := os.OpenFile(ioRdLimitPath, os.O_WRONLY|os.O_CREATE, 0)
	if err != nil {
		return nil, err
	}
	defer fr.Close()
	fw, err := os.OpenFile(ioWrLimitPath, os.O_WRONLY|os.O_CREATE, 0)
	if err != nil {
		return nil, err
	}
	defer fw.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,
		}

		rlimitStr := fmt.Sprintf("%s:%s %d", devinfo.Major, devinfo.Minor, rlimit)
		wlimitStr := fmt.Sprintf("%s:%s %d", devinfo.Major, devinfo.Minor, wlimit)

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

		// set io write limit
		_, err = fw.WriteString(wlimitStr)
		if err != nil {
			l.Debugf("set %s io write limit error %v, config:%s", devinfo.Name, err, wlimitStr)
		} else {
			devLmts[devKey].WrLimit = true
			ioWrLimitSuccCnt++
			l.Debugf("set %s io write limit success, config:%s", devinfo.Name, wlimitStr)
		}
	}
	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 blkio.throttle.io_service_bytes
// 8:0 Read 5829550080
// 8:0 Write 4474982400
// 8:0 Sync 4415606784
// 8:0 Async 5888925696
// 8:0 Total 10304532480
// Total 10304532480
func GetAllDevIOStat(cgroupDir string) (map[string]*types.IOStat, error) {
	data, err := os.ReadFile(filepath.Join(cgroupDir, "blkio.throttle.io_service_bytes"))
	if err != nil {
		return nil, err
	}

	ioStats := make(map[string]*types.IOStat, 0)
	ioStats["total"] = &types.IOStat{}
	for _, line := range strings.Split(string(data), "\n") {
		line = strings.TrimSpace(line)
		// line content: 8:0 Read 8433958912
		fields := strings.Split(line, " ")
		if len(fields) != 3 {
			continue
		}
		device := fields[0]
		event := fields[1]
		size := fields[2]

		if event == "Read" || event == "Write" {
			sizeUint64, err := strconv.ParseUint(size, 10, 64)
			if err != nil {
				return ioStats, fmt.Errorf("unable to convert size to uint64: %s", err)
			}

			if _, ok := ioStats[device]; !ok {
				ioStats[device] = &types.IOStat{}
			}
			if event == "Read" {
				ioStats[device].Read = sizeUint64
				ioStats["total"].Read += sizeUint64
			}
			if event == "Write" {
				ioStats[device].Write = sizeUint64
				ioStats["total"].Write += sizeUint64
			}
		}
	}
	return ioStats, nil
}