package resources

import (
	"sort"
	"time"

	"github.com/shirou/gopsutil/disk"
)

type HostDiskIOUsage struct {
	Device     string
	ReadBytes  uint64
	WriteBytes uint64
	ReadUsage  uint64
	WriteUsage uint64
}

type HostDiskIOUsageList []HostDiskIOUsage

func (e HostDiskIOUsageList) Len() int {
	return len(e)
}

func (e HostDiskIOUsageList) Less(i, j int) bool {
	return e[i].Device < e[j].Device
}

func (e HostDiskIOUsageList) Swap(i, j int) {
	e[i], e[j] = e[j], e[i]
}

type HostDiskIOCountersStatWrap struct {
	IOStat map[string]disk.IOCountersStat
	Ts     int64
}

func GetHostDiskIOUsageWithInterval(d time.Duration) ([]HostDiskIOUsage, error) {
	usages := make([]HostDiskIOUsage, 0)

	pre, err := disk.IOCounters()
	if err != nil {
		return usages, err
	}
	preWrap := HostDiskIOCountersStatWrap{IOStat: pre, Ts: time.Now().Unix()}

	time.Sleep(d)

	now, err := disk.IOCounters()
	if err != nil {
		return usages, err
	}
	nowWrap := HostDiskIOCountersStatWrap{IOStat: now, Ts: time.Now().Unix()}

	return calculateDiskIOUsage(&preWrap, &nowWrap), nil
}

func GetHostDiskIOUsageWithPre(preWrap *HostDiskIOCountersStatWrap) (
	[]HostDiskIOUsage, *HostDiskIOCountersStatWrap, error) {
	usages := make([]HostDiskIOUsage, 0)

	now, err := disk.IOCounters()
	if err != nil {
		return usages, nil, err
	}
	nowWrap := HostDiskIOCountersStatWrap{IOStat: now, Ts: time.Now().Unix()}

	if preWrap == nil {
		return calculateDiskIOUsage(&nowWrap, &nowWrap), &nowWrap, nil
	}

	return calculateDiskIOUsage(preWrap, &nowWrap), &nowWrap, nil
}

func calculateDiskIOUsage(preWrap, nowWrap *HostDiskIOCountersStatWrap) []HostDiskIOUsage {
	usages := make([]HostDiskIOUsage, 0)

	duration := nowWrap.Ts - preWrap.Ts
	for dev, nowStat := range nowWrap.IOStat {
		usage := HostDiskIOUsage{
			Device:     dev,
			ReadBytes:  nowStat.ReadBytes,
			WriteBytes: nowStat.WriteBytes,
		}

		if preStat, ok := preWrap.IOStat[dev]; ok {
			rdiff := (nowStat.ReadBytes - preStat.ReadBytes)
			if rdiff > 0 && duration > 0 {
				usage.ReadUsage = rdiff / uint64(duration)
			}

			wdiff := (nowStat.WriteBytes - preStat.WriteBytes)
			if wdiff > 0 && duration > 0 {
				usage.WriteUsage = wdiff / uint64(duration)
			}
		}

		usages = append(usages, usage)
	}

	sort.Sort(HostDiskIOUsageList(usages))
	return usages
}