package backend

import (
	"get-container/gpu"
	"sync"
)

/*
	backend包为tui包提供数据
*/

var (
	MapIdDCU = sync.Map{} // 记录dcu信息
)

type DCUInfo struct {
	Id               int
	Name             string // full
	PerformanceLevel string // full
	Fan              string // full
	Temp             float32
	PwrAvg           float32
	PwrCap           float32 // full
	BusId            string  // full
	MemTotal         int
	MemUsed          int
	MemUsedPerent    float32
	Mig              bool // full
	DCUUTil          float32
	Ecc              bool // full
	PwrMode          string
}

// UpdateDCUInfo 更新dcu信息，full表示是否全部更新
func UpdateDCUInfo(full bool) {
	wg := sync.WaitGroup{}
	var smiAll map[int]*gpu.SMIAllOutput
	var eccInfo, migInfo map[int]bool
	var memInfo map[int]gpu.DCUMemInfo
	var runInfo map[int]gpu.DCURunningInfo
	var errSmiAll, errEcc, errMem, errRun error
	if full {
		wg.Add(5)
		go func() {
			smiAll, errSmiAll = gpu.GetSMIAllOutput()
			wg.Done()
		}()
		go func() {
			eccInfo, errEcc = gpu.GetEccInfo()
			wg.Done()
		}()
	} else {
		wg.Add(3)
	}
	go func() {
		migInfo = gpu.GetMigInfo()
		wg.Done()
	}()
	go func() {
		memInfo, errMem = gpu.GetDCUMemInfo()
		wg.Done()
	}()
	go func() {
		runInfo, errRun = gpu.GetRunningInfo()
		wg.Done()
	}()
	wg.Wait()
	cache := make(map[int]DCUInfo)
	if errSmiAll == nil && smiAll != nil {
		for k, v := range smiAll {
			i, have := cache[k]
			if !have {
				i = DCUInfo{}
				i.Id = k
			}
			i.Name = v.CardSeries
			i.PerformanceLevel = v.PerLevel
			i.Fan = gpu.NA
			i.PwrCap = v.MaxPwr
			i.PwrAvg = v.AvgPwr
			i.BusId = v.PCIBus
			i.MemUsedPerent = v.HCUMemUsage
			i.DCUUTil = v.HCUUsage
			cache[k] = i
		}
	}
	if errEcc == nil && eccInfo != nil {
		for k, v := range eccInfo {
			i, have := cache[k]
			if !have {
				i = DCUInfo{}
				i.Id = k
			}
			i.Ecc = v
			cache[k] = i
		}
	}
	for k, v := range migInfo {
		i, have := cache[k]
		if !have {
			i = DCUInfo{}
			i.Id = k
		}
		i.Mig = v
		cache[k] = i
	}
	if memInfo != nil && errMem == nil {
		for k, v := range memInfo {
			i, have := cache[k]
			if !have {
				i = DCUInfo{}
				i.Id = k
			}
			i.MemTotal = int(v.Total.Num)
			i.MemUsed = int(v.Used.Num)
			cache[k] = i
		}
	}
	if errRun == nil && runInfo != nil {
		for k, v := range runInfo {
			i, have := cache[k]
			if !have {
				i = DCUInfo{}
				i.Id = k
			}
			i.Temp = v.Temp
			i.PwrAvg = v.AvgPower
			i.PerformanceLevel = v.PerformanceLevel
			i.MemUsedPerent = v.MemPerc
			i.DCUUTil = v.HCUPerc
			i.PwrMode = v.Mode
			cache[k] = i
		}
	}
	for k, v := range cache {
		old, have := MapIdDCU.LoadOrStore(k, &v)
		if !have {
			continue
		}
		d := old.(*DCUInfo)
		if full {
			d.Name = v.Name
			d.PerformanceLevel = v.PerformanceLevel
			d.Fan = v.Fan
			d.PwrCap = v.PwrCap
			d.BusId = v.BusId
			d.Mig = v.Mig
			d.Ecc = v.Ecc
		}
		d.Temp = v.Temp
		d.PwrAvg = v.PwrAvg
		d.MemTotal = v.MemTotal
		d.MemUsed = v.MemUsed
		d.MemUsedPerent = v.MemUsedPerent
		d.DCUUTil = v.DCUUTil
		d.PwrMode = v.PwrMode
	}
}

func GetDCUInfo() map[int]DCUInfo {
	result := make(map[int]DCUInfo)
	MapIdDCU.Range(func(key, value any) bool {
		id := key.(int)
		val := value.(*DCUInfo)
		result[id] = *val
		return true
	})
	return result
}