package collector

import (
	"bee/audit/internal/schema"
	"fmt"
	"testing"
)

func TestParseMstflintQuery(t *testing.T) {
	raw := `Device #1:
----------
FW Version:     28.39.1002
Board Serial Number: MT1234ABC
`
	fw, serial := parseMstflintQuery(raw)
	if fw != "28.39.1002" {
		t.Fatalf("firmware: got %q", fw)
	}
	if serial != "MT1234ABC" {
		t.Fatalf("serial: got %q", serial)
	}
}

func TestParseEthtoolFirmwareInfo(t *testing.T) {
	raw := `driver: mlx5_core
version: 6.6.31-0-lts
firmware-version: 28.39.1002 (MT_0000000000)
bus-info: 0000:18:00.0
`
	fw := parseEthtoolFirmwareInfo(raw)
	if fw != "28.39.1002 (MT_0000000000)" {
		t.Fatalf("firmware: got %q", fw)
	}
}

func TestEnrichPCIeWithMellanox_mstflint(t *testing.T) {
	origMst := mstflintQuery
	origEth := ethtoolInfoQuery
	origIfaces := netIfacesByBDF
	t.Cleanup(func() {
		mstflintQuery = origMst
		ethtoolInfoQuery = origEth
		netIfacesByBDF = origIfaces
	})

	mstflintQuery = func(bdf string) (string, error) {
		if bdf != "0000:18:00.0" {
			t.Fatalf("unexpected bdf: %s", bdf)
		}
		return "FW Version: 28.39.1002\nBoard Serial Number: SN-MST-001\n", nil
	}
	ethtoolInfoQuery = func(string) (string, error) {
		t.Fatal("ethtool should not be called when mstflint succeeds")
		return "", nil
	}
	netIfacesByBDF = func(string) []string { return nil }

	vendorID := mellanoxVendorID
	bdf := "0000:18:00.0"
	manufacturer := "Mellanox Technologies"
	devs := []schema.HardwarePCIeDevice{{
		VendorID:     &vendorID,
		BDF:          &bdf,
		Manufacturer: &manufacturer,
	}}

	out := enrichPCIeWithMellanox(devs)
	if out[0].Firmware == nil || *out[0].Firmware != "28.39.1002" {
		t.Fatalf("firmware: got %v", out[0].Firmware)
	}
	if out[0].SerialNumber == nil || *out[0].SerialNumber != "SN-MST-001" {
		t.Fatalf("serial: got %v", out[0].SerialNumber)
	}
}

func TestEnrichPCIeWithMellanox_fallbackEthtool(t *testing.T) {
	origMst := mstflintQuery
	origEth := ethtoolInfoQuery
	origIfaces := netIfacesByBDF
	t.Cleanup(func() {
		mstflintQuery = origMst
		ethtoolInfoQuery = origEth
		netIfacesByBDF = origIfaces
	})

	mstflintQuery = func(string) (string, error) {
		return "", fmt.Errorf("mstflint not found")
	}
	netIfacesByBDF = func(bdf string) []string {
		if bdf != "0000:18:00.0" {
			t.Fatalf("unexpected bdf: %s", bdf)
		}
		return []string{"eth0"}
	}
	ethtoolInfoQuery = func(iface string) (string, error) {
		if iface != "eth0" {
			t.Fatalf("unexpected iface: %s", iface)
		}
		return "driver: mlx5_core\nfirmware-version: 28.40.1000\n", nil
	}

	vendorID := mellanoxVendorID
	bdf := "0000:18:00.0"
	manufacturer := "NVIDIA Networking"
	devs := []schema.HardwarePCIeDevice{{
		VendorID:     &vendorID,
		BDF:          &bdf,
		Manufacturer: &manufacturer,
	}}

	out := enrichPCIeWithMellanox(devs)
	if out[0].Firmware == nil || *out[0].Firmware != "28.40.1000" {
		t.Fatalf("firmware: got %v", out[0].Firmware)
	}
	if out[0].SerialNumber != nil {
		t.Fatalf("serial should stay nil without mstflint, got %v", out[0].SerialNumber)
	}
}