package collector

import (
	"bee/audit/internal/schema"
	"log/slog"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

var (
	queryPCILSPCIDetail = func(bdf string) (string, error) {
		out, err := exec.Command("lspci", "-vv", "-s", bdf).Output()
		if err != nil {
			return "", err
		}
		return string(out), nil
	}
	readPCIVPDFile = func(bdf string) ([]byte, error) {
		return os.ReadFile(filepath.Join("/sys/bus/pci/devices", bdf, "vpd"))
	}
)

func enrichPCIeWithPCISerials(devs []schema.HardwarePCIeDevice) []schema.HardwarePCIeDevice {
	enriched := 0
	for i := range devs {
		if !shouldProbePCIeSerial(devs[i]) {
			continue
		}
		bdf := normalizePCIeBDF(*devs[i].BDF)
		if bdf == "" {
			continue
		}
		if serial := queryPCIDeviceSerial(bdf); serial != "" {
			devs[i].SerialNumber = &serial
			enriched++
		}
	}
	if enriched > 0 {
		slog.Info("pcie: serials enriched", "count", enriched)
	}
	return devs
}

func shouldProbePCIeSerial(dev schema.HardwarePCIeDevice) bool {
	if dev.BDF == nil || dev.SerialNumber != nil {
		return false
	}
	if dev.DeviceClass == nil {
		return false
	}
	class := strings.TrimSpace(*dev.DeviceClass)
	return isNICClass(class) || isGPUClass(class)
}

func queryPCIDeviceSerial(bdf string) string {
	if out, err := queryPCILSPCIDetail(bdf); err == nil {
		if serial := parseLSPCIDetailSerial(out); serial != "" {
			return serial
		}
	}
	if raw, err := readPCIVPDFile(bdf); err == nil {
		return parsePCIVPDSerial(raw)
	}
	return ""
}

func parseLSPCIDetailSerial(raw string) string {
	for _, line := range strings.Split(raw, "\n") {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		lower := strings.ToLower(line)
		if !strings.Contains(lower, "serial number:") {
			continue
		}
		idx := strings.Index(line, ":")
		if idx < 0 {
			continue
		}
		if serial := strings.TrimSpace(line[idx+1:]); serial != "" {
			return serial
		}
	}
	return ""
}

func parsePCIVPDSerial(raw []byte) string {
	for i := 0; i+3 < len(raw); i++ {
		if raw[i] != 'S' || raw[i+1] != 'N' {
			continue
		}
		length := int(raw[i+2])
		if length <= 0 || length > 64 || i+3+length > len(raw) {
			continue
		}
		value := strings.TrimSpace(strings.Trim(string(raw[i+3:i+3+length]), "\x00"))
		if !looksLikeSerial(value) {
			continue
		}
		return value
	}
	return ""
}

func looksLikeSerial(value string) bool {
	if len(value) < 4 {
		return false
	}
	hasAlphaNum := false
	for _, r := range value {
		switch {
		case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9':
			hasAlphaNum = true
		case strings.ContainsRune(" -_./:", r):
		default:
			return false
		}
	}
	return hasAlphaNum
}