package gpu

import (
	"encoding/json"
	"os"
	"sort"
	"strconv"
	"strings"
	"testing"
	"time"
)

const (
	Data1 = `============================ System Management Interface =============================
======================================================================================
No KFD PIDs currently running!
======================================================================================
=================================== End of SMI Log ===================================`
	Data2 = `================================= System Management Interface ==================================
================================================================================================
PIDs for KFD processes:

PID: 142211
	   PASID: 32775
	   HCU Node(Include CPU sort): ['2']
	   HCU Index: ['0']
	   GPUID: ['15868']
	   PCI BUS: ['0000:49:00.0']
	   VRAM USED(MiB): 54682
	   VRAM USED(%): 83
	   SDMA USED: 0

PID: 142218
	   PASID: 32772
	   HCU Node(Include CPU sort): ['8']
	   HCU Index: ['6']
	   GPUID: ['37441']
	   PCI BUS: ['0000:cd:00.0']
	   VRAM USED(MiB): 43485
	   VRAM USED(%): 66
	   SDMA USED: 0

PID: 142236
	   PASID: 32774
	   HCU Node(Include CPU sort):
	   HCU Index:
	   GPUID:
	   PCI BUS:
	   VRAM USED(MiB): 0
	   VRAM USED(%): inf
	   SDMA USED: 0

PID: 142216
	   PASID: 32773
	   HCU Node(Include CPU sort): ['6']
	   HCU Index: ['4']
	   GPUID: ['18915']
	   PCI BUS: ['0000:9c:00.0']
	   VRAM USED(MiB): 37438
	   VRAM USED(%): 57
	   SDMA USED: 0

PID: 142214
	   PASID: 32771
	   HCU Node(Include CPU sort): ['3']
	   HCU Index: ['1']
	   GPUID: ['51742']
	   PCI BUS: ['0000:54:00.0']
	   VRAM USED(MiB): 54815
	   VRAM USED(%): 84
	   SDMA USED: 0

PID: 142212
	   PASID: 32768
	   HCU Node(Include CPU sort): ['4']
	   HCU Index: ['2']
	   GPUID: ['14451']
	   PCI BUS: ['0000:5e:00.0']
	   VRAM USED(MiB): 40722
	   VRAM USED(%): 62
	   SDMA USED: 0

PID: 142249
	   PASID: 32780
	   HCU Node(Include CPU sort):
	   HCU Index:
	   GPUID:
	   PCI BUS:
	   VRAM USED(MiB): 0
	   VRAM USED(%): inf
	   SDMA USED: 0

PID: 142239
	   PASID: 32776
	   HCU Node(Include CPU sort):
	   HCU Index:
	   GPUID:
	   PCI BUS:
	   VRAM USED(MiB): 0
	   VRAM USED(%): inf
	   SDMA USED: 0

PID: 142247
	   PASID: 32781
	   HCU Node(Include CPU sort):
	   HCU Index:
	   GPUID:
	   PCI BUS:
	   VRAM USED(MiB): 0
	   VRAM USED(%): inf
	   SDMA USED: 0

PID: 142217
	   PASID: 32770
	   HCU Node(Include CPU sort): ['5']
	   HCU Index: ['3']
	   GPUID: ['34940']
	   PCI BUS: ['0000:67:00.0']
	   VRAM USED(MiB): 40826
	   VRAM USED(%): 62
	   SDMA USED: 0

PID: 142245
	   PASID: 32779
	   HCU Node(Include CPU sort):
	   HCU Index:
	   GPUID:
	   PCI BUS:
	   VRAM USED(MiB): 0
	   VRAM USED(%): inf
	   SDMA USED: 0

PID: 142235
	   PASID: 32784
	   HCU Node(Include CPU sort):
	   HCU Index:
	   GPUID:
	   PCI BUS:
	   VRAM USED(MiB): 0
	   VRAM USED(%): inf
	   SDMA USED: 0

PID: 142215
	   PASID: 32797
	   HCU Node(Include CPU sort): ['9']
	   HCU Index: ['7']
	   GPUID: ['46537']
	   PCI BUS: ['0000:dd:00.0']
	   VRAM USED(MiB): 43279
	   VRAM USED(%): 66
	   SDMA USED: 0

PID: 142243
	   PASID: 32778
	   HCU Node(Include CPU sort):
	   HCU Index:
	   GPUID:
	   PCI BUS:
	   VRAM USED(MiB): 0
	   VRAM USED(%): inf
	   SDMA USED: 0

PID: 142213
	   PASID: 32769
	   HCU Node(Include CPU sort): ['7']
	   HCU Index: ['5']
	   GPUID: ['4240']
	   PCI BUS: ['0000:bc:00.0']
	   VRAM USED(MiB): 37436
	   VRAM USED(%): 57
	   SDMA USED: 0

PID: 142241
	   PASID: 32777
	   HCU Node(Include CPU sort):
	   HCU Index:
	   GPUID:
	   PCI BUS:
	   VRAM USED(MiB): 0
	   VRAM USED(%): inf
	   SDMA USED: 0

================================================================================================
======================================== End of SMI Log ========================================`
)

func TestRegexp(t *testing.T) {
	str := "[\"7\"]"
	ss := make([]string, 0)
	err := json.Unmarshal([]byte(str), &ss)
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("%v", ss)
}

func TestParseDCUPidInfo(t *testing.T) {
	i, e := parseDCUPidInfo(Data2)
	if e != nil {
		t.Fatal(e)
	}
	for _, info := range i {
		t.Logf("%+v\n", info)
	}

	ii, e := parseDCUPidInfo(Data1)
	if e != nil {
		t.Fatal(e)
	}
	for _, info := range ii {
		t.Logf("%+v\n", info)
	}
}

const DCUPidInfoStr = `
============================ System Management Interface =============================
======================================================================================
HCU     Temp     AvgPwr     Perf     PwrCap     VRAM%      HCU%      Mode
0       34.0C    140.0W     manual   800.0W     0%         0.0%      Normal
1       35.0C    140.0W     manual   800.0W     0%         0.0%      Normal
2       35.0C    140.0W     manual   800.0W     0%         0.0%      Normal
3       33.0C    140.0W     manual   800.0W     0%         0.0%      Normal
4       36.0C    140.0W     manual   800.0W     0%         0.0%      Normal
5       36.0C    140.0W     manual   800.0W     0%         0.0%      Normal
6       34.0C    140.0W     manual   800.0W     0%         0.0%      Normal
7       34.0C    140.0W     manual   800.0W     0%         0.0%      Normal
======================================================================================
=================================== End of SMI Log ===================================

`

func TestParseRunningInfo(t *testing.T) {
	bytes, err := os.ReadFile("../test-data/hy-smi-1.17.0-1.log")
	if err != nil {
		t.Error(err)
	}
	str := strings.TrimSpace(strings.Trim(string(bytes), "\b"))
	i, e := parseRunningInfo(str)
	if e != nil {
		t.Fatal(e)
	}
	for _, info := range i {
		t.Logf("%+v\n", info)
	}
}

func TestParseSMIAllOutput(t *testing.T) {
	b, e := os.ReadFile("../test-data/hy-smi-1.17.0-a.log")
	if e != nil {
		t.Error(e)
	}
	timeStart := time.Now()
	lines := strings.Split(strings.Trim(strings.TrimSpace(string(b)), "\n"), "\n")
	info := make(map[int][]string)
	for _, line := range lines {
		if ReUselessLine.MatchString(line) || ReEmptyLine.MatchString(line) {
			continue
		}
		if ReInfoHeader.MatchString(line) {
			fields := ReInfoHeader.FindStringSubmatch(strings.TrimSpace(strings.ReplaceAll(line, "\t", " ")))
			if len(fields) <= 2 {
				continue
			}
			id, innerErr := strconv.Atoi(fields[1])
			if innerErr != nil {
				t.Error(innerErr)
			}
			if v, ok := info[id]; !ok {
				info[id] = make([]string, 0)
				info[id] = append(v, fields[2])
			} else {
				info[id] = append(v, fields[2])
			}
		}
	}
	result := make([]*SMIAllOutput, 0)
	for k, v := range info {
		item, innerErr := parseSMIAllOutput(k, strings.Join(v, "\n"))
		if innerErr != nil {
			t.Error(innerErr)
		}
		if item != nil {
			result = append(result, item)
		}
	}
	sort.Slice(result, func(i, j int) bool {
		return result[i].Id < result[j].Id
	})
	end := time.Now()
	t.Log(end.Sub(timeStart).Seconds())
	for _, i := range result {
		t.Logf("%+v", *i)
	}
}

func TestPCIBus(t *testing.T) {
	f := RePCIBus.FindStringSubmatch("PCI Bus: 0000:9f:00.0 --> Socket id: 1")
	for _, l := range f {
		t.Logf("%s", l)
	}

	if s := regMatch(RePCIBus, "PCI Bus: 0000:9f:00.0 --> Socket id: 1", 1); s != nil {
		t.Logf("%+v", s)
	}

}
