package gogpu

/*
#cgo CFLAGS: -I. -w -g
#cgo CXXFLAGS: -I. -std=c++11 -w -g
#cgo LDFLAGS: -ldl
#include <stdlib.h>
#include "cgpu.h"
*/
import "C"
import (
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strconv"
	"strings"
	"unsafe"
)

// ValidGPU pass needed gpu memory size test if satisfy
func ValidGPU(memSize int) int {
	return int(C.get_idle_gpu(C.int(memSize)))
}

// SatisfyGPU satisfy unit "M"
func SatisfyGPU(index, memSize, reserve int) bool {
	info, err := Info()
	if err != nil || info.Count >= index {
		fmt.Println("SatisfyGPU no gpu or index ilegal: ", index, " gpu count : ", info.Count)
		return false
	}

	var M1 int64 = 1024 * 1024
	if info.Info[index].GpuMemoryFree-int64(memSize)*M1 > int64(reserve)*M1 {
		return true
	}
	return false
}

// GPURank sort
type GPURank []GpuUnitInfo

func (a GPURank) Len() int           { return len(a) }
func (a GPURank) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a GPURank) Less(i, j int) bool { return a[i].GpuMemoryFree < a[j].GpuMemoryFree }

// RankGPU rank
func RankGPU() []int {
	info, err := Info()
	if err != nil {
		fmt.Println("RankGPU no gpu")
		return nil
	}
	var ret []int
	if info.Count == 1 {
		ret = append(ret, 0)
	} else {
		sort.Sort(GPURank(info.Info))
		for i := len(info.Info) - 1; i >= 0; i-- {
			ret = append(ret, info.Info[i].GpuIndex)
		}
	}
	return ret
}

// IdleGPU idle most
func IdleGPU() int {
	info, err := Info()
	if err != nil {
		fmt.Println("IdleGPU no gpu")
		return -1
	}
	if info.Count == 1 {
		return 0
	}

	var free int64
	gpu := 0
	for k, v := range info.Info {
		if v.GpuMemoryFree > free {
			free = v.GpuMemoryFree
			gpu = k
		}
	}
	return gpu
}

// GpuUnitInfo gpu card info
type GpuUnitInfo struct {
	GpuIndex       int   `json:"index"`
	GpuUtilization int   `json:"utilization"`
	GpuTemperature int   `json:"temperature"`
	GpuDecoderUtil int   `json:"decoder"`
	GpuEncoderUtil int   `json:"encoder"`
	GpuProcesses   int   `json:"process"`
	GpuMemoryTotal int64 `json:"memory"`
	GpuMemoryFree  int64 `json:"free"`
	GpuMemoryUsed  int64 `json:"used"`
}

func (u GpuUnitInfo) String() string {
	s, _ := json.Marshal(u)
	return string(s)
}

// GpuInfo gpu info
type GpuInfo struct {
	Count int           `json:"count"`
	Info  []GpuUnitInfo `json:"info"`
}

func (g GpuInfo) String() string {
	s, _ := json.Marshal(g)
	return string(s)
}

// Info gpu infos
func Info() (*GpuInfo, error) {

	p := C.get_gpu_info()
	if p == nil {
		return nil, errors.New("can't get gpu infos")
	}
	defer C.free(unsafe.Pointer(p))

	str := C.GoString(p)

	infos := strings.Split(str, "|")

	length := len(infos)

	var units []GpuUnitInfo
	var count int
	for i := 0; i < length; i++ {
		if infos[i] == "gpu" {
			u := GpuUnitInfo{}
			u.GpuIndex = count
			count++

			i++
			u.GpuUtilization, _ = strconv.Atoi(infos[i])
			i++
			u.GpuTemperature, _ = strconv.Atoi(infos[i])
			i++
			u.GpuDecoderUtil, _ = strconv.Atoi(infos[i])
			i++
			u.GpuEncoderUtil, _ = strconv.Atoi(infos[i])
			i++
			u.GpuProcesses, _ = strconv.Atoi(infos[i])
			i++
			u.GpuMemoryTotal, _ = strconv.ParseInt(infos[i], 10, 64)
			i++
			u.GpuMemoryFree, _ = strconv.ParseInt(infos[i], 10, 64)
			i++
			u.GpuMemoryUsed, _ = strconv.ParseInt(infos[i], 10, 64)
			units = append(units, u)
		}
	}

	return &GpuInfo{count, units}, nil
}