package main

import (
	"fmt"
	"io/ioutil"
	"path"
	"path/filepath"
	"strings"

	"golang.fcuny.net/numap/internal/hwids"
	"golang.fcuny.net/numap/internal/sysfs"
)

const (
	node_root      = "/sys/devices/system/node/node*"
	CLASS_NVMe     = 67586
	CLASS_ETHERNET = 131072
	CLASS_GPU      = 197120
)

type node struct {
	Name       string      `json:"name"`
	Path       string      `json:"path"`
	CpuList    string      `json:"cpulist"`
	PCIDevices []PCIDevice `json:"pci_devices"`
}

type PCIDevice struct {
	Vendor string `json:"vendor"`
	Name   string `json:"name"`
}

func findNodes(hwdb hwids.PciDevices) (map[string]node, error) {
	nodes := make(map[string]node)

	files, err := filepath.Glob(node_root)
	if err != nil {
		return nil, fmt.Errorf("Failed to find NUMA nodes under %s: %+v", node_root, err)
	}
	if len(files) == 0 {
		return nil, fmt.Errorf("Could not find NUMA node in %s", node_root)
	}

	for _, f := range files {
		n, err := newNode(f)
		if err != nil {
			return make(map[string]node), err
		}
		nodes[n.Name] = n
	}

	r, err := mapPCIDevicesToNumaNode(hwdb)
	if err != nil {
		panic(err)
	}
	for k, v := range r {
		nodeName := fmt.Sprintf("node%d", k)
		n := nodes[nodeName]
		n.PCIDevices = v
		nodes[nodeName] = n
	}
	return nodes, nil
}

func mapPCIDevicesToNumaNode(hwdb hwids.PciDevices) (map[int][]PCIDevice, error) {
	devices := sysfs.ScanPCIDevices()
	r := map[int][]PCIDevice{}

	for _, d := range devices {
		if d.Class == CLASS_NVMe || d.Class == CLASS_ETHERNET || d.Class == CLASS_GPU {
			_, ok := hwdb[uint16(d.Vendor)]
			if ok {
				desc := hwdb[uint16(d.Vendor)]
				var vendor, name string
				for _, m := range desc {
					if uint64(m.Device) == d.Device && uint64(m.Vendor) == d.Vendor {
						vendor = m.VendorName
						name = m.DeviceName
						break
					}
				}
				pciDevice := PCIDevice{
					Vendor: vendor,
					Name:   name,
				}
				r[d.NumaNode] = append(r[d.NumaNode], pciDevice)
			}
		}
	}
	return r, nil
}

func newNode(p string) (node, error) {
	_, name := path.Split(p)

	cpulist, err := cpuList(p)
	if err != nil {
		return node{}, err
	}

	return node{
		Name:       name,
		Path:       p,
		CpuList:    cpulist,
		PCIDevices: []PCIDevice{},
	}, nil
}

func cpuList(p string) (string, error) {
	lpath := filepath.Join(p, "cpulist")
	c, err := ioutil.ReadFile(lpath)
	if err != nil {
		return "", fmt.Errorf("Failed to open %s: %+v", lpath, err)
	}
	return strings.TrimRight(string(c), "\n"), nil
}