package main

import (
	"bufio"
	"fmt"
	"os"
	"strings"

	"golang.org/x/sys/windows"
	"golang.org/x/sys/windows/registry"
)

// 显卡类 GUID
var gpuClassGUID = windows.GUID{
	Data1: 0x4d36e968,
	Data2: 0xe325,
	Data3: 0x11ce,
	Data4: [8]byte{0xbf, 0xc1, 0x08, 0x00, 0x2b, 0xe1, 0x03, 0x18},
}

// GPUInfo 保存一张显卡的信息
type GPUInfo struct {
	Name string // 显示名称
	Path string // 注册表枚举路径（HKLM\SYSTEM\CurrentControlSet\Enum\...）
	ID   string // 设备实例 ID
}

func main() {
	fmt.Println("============================================")
	fmt.Println("        显卡名称修改工具 (GPU Rename)        ")
	fmt.Println("============================================")
	fmt.Println()

	// 1. 枚举当前实际存在的物理显卡
	gpus, err := enumerateGPUs()
	if err != nil {
		fmt.Printf("错误：枚举显卡失败 - %v\n", err)
		fmt.Println("提示：请以管理员身份运行本程序。")
		waitEnter()
		return
	}

	if len(gpus) == 0 {
		fmt.Println("错误：未检测到任何物理显卡设备。")
		waitEnter()
		return
	}

	// 2. 选择目标显卡
	var target GPUInfo
	if len(gpus) > 1 {
		fmt.Printf("检测到 %d 块物理显卡，请选择要修改的显卡：\n", len(gpus))
		for i, g := range gpus {
			fmt.Printf("  [%d] %s\n", i+1, g.Name)
		}
		fmt.Println()
		target, err = chooseGPU(gpus)
		if err != nil {
			fmt.Printf("错误：%v\n", err)
			waitEnter()
			return
		}
	} else {
		target = gpus[0]
		fmt.Printf("检测到 1 块物理显卡：%s\n\n", target.Name)
	}

	fmt.Printf("当前选中显卡：%s\n", target.Name)
	fmt.Printf("设备实例 ID：%s\n\n", target.ID)

	// 3. 输入新名称
	newName := readNewName()
	if newName == "" {
		fmt.Println("已取消操作。")
		waitEnter()
		return
	}

	fmt.Printf("即将把显卡名称修改为：%s\n", newName)

	// 4. 执行修改
	if err := renameGPU(target, newName); err != nil {
		fmt.Printf("错误：修改失败 - %v\n", err)
		fmt.Println("请确认已使用管理员权限运行。")
		waitEnter()
		return
	}

	fmt.Println("修改已写入注册表，正在验证...")

	// 5. 验证修改结果
	verifyName, err := verifyGPU(target)
	if err != nil {
		fmt.Printf("警告：验证时发生错误 - %v\n", err)
	} else if verifyName == newName {
		fmt.Println("验证通过：显卡名称已成功修改为指定值。")
	} else {
		fmt.Printf("警告：验证结果与预期不一致，当前值为：%s\n", verifyName)
		fmt.Println("      部分设备可能需要重启后才会在设备管理器中刷新。")
	}

	fmt.Println()
	fmt.Println("============================================")
	fmt.Println("  修改成功！请重启机器以使修改完全生效。   ")
	fmt.Println("============================================")

	waitEnter()
}

// enumerateGPUs 使用 SetupAPI 枚举当前实际存在的物理显卡。
// 使用 DIGCF_PRESENT 标志，只列出真正连接在机器上的设备，
// 并过滤掉虚拟显示适配器（无物理设备对象名）。
func enumerateGPUs() ([]GPUInfo, error) {
	// 只枚举当前存在的设备（DIGCF_PRESENT）
	devSet, err := windows.SetupDiGetClassDevsEx(
		&gpuClassGUID,
		"",
		0,
		windows.DIGCF_PRESENT,
		0,
		"",
	)
	if err != nil {
		return nil, err
	}
	defer windows.SetupDiDestroyDeviceInfoList(devSet)

	var gpus []GPUInfo
	for i := 0; ; i++ {
		devInfo, err := windows.SetupDiEnumDeviceInfo(devSet, i)
		if err != nil {
			break // 枚举完毕
		}

		// 过滤：只保留物理显卡（PCI 总线上的设备），
		// 排除 ROOT\SWD 等虚拟显示适配器。
		enumerator := ""
		if v, err := windows.SetupDiGetDeviceRegistryProperty(devSet, devInfo, windows.SPDRP_ENUMERATOR_NAME); err == nil {
			if s, ok := v.(string); ok {
				enumerator = s
			}
		}
		if !strings.EqualFold(enumerator, "PCI") {
			continue
		}

		// 获取实例 ID
		instanceID, err := windows.SetupDiGetDeviceInstanceId(devSet, devInfo)
		if err != nil {
			continue
		}

		// 获取显示名称（优先 FriendlyName，其次 DeviceDesc）
		name := getDeviceName(devSet, devInfo, instanceID)

		gpus = append(gpus, GPUInfo{
			Name: name,
			Path: instanceIDToRegPath(instanceID),
			ID:   instanceID,
		})
	}

	return gpus, nil
}

// getDeviceName 获取设备显示名称，优先 FriendlyName，回退到 DeviceDesc
func getDeviceName(devSet windows.DevInfo, devInfo *windows.DevInfoData, fallback string) string {
	if v, err := windows.SetupDiGetDeviceRegistryProperty(devSet, devInfo, windows.SPDRP_FRIENDLYNAME); err == nil {
		if s, ok := v.(string); ok && strings.TrimSpace(s) != "" {
			return s
		}
	}
	if v, err := windows.SetupDiGetDeviceRegistryProperty(devSet, devInfo, windows.SPDRP_DEVICEDESC); err == nil {
		if s, ok := v.(string); ok && strings.TrimSpace(s) != "" {
			return cleanDeviceDesc(s)
		}
	}
	return fallback
}

// instanceIDToRegPath 将设备实例 ID 转换为注册表枚举路径。
// 例如 PCI\VEN_10DE&DEV_1DB1...\4&31024D2F&0&0018
//      -> SYSTEM\CurrentControlSet\Enum\PCI\VEN_10DE&DEV_1DB1...\4&31024D2F&0&0018
func instanceIDToRegPath(instanceID string) string {
	return `SYSTEM\CurrentControlSet\Enum\` + instanceID
}

// cleanDeviceDesc 去掉 DeviceDesc 中的 INF 引用前缀，只保留可读名称
func cleanDeviceDesc(desc string) string {
	if idx := strings.LastIndex(desc, ";"); idx >= 0 && idx < len(desc)-1 {
		return desc[idx+1:]
	}
	return desc
}

// getDeviceDesc 读取注册表中的 DeviceDesc
func getDeviceDesc(regPath string) (string, error) {
	k, err := registry.OpenKey(registry.LOCAL_MACHINE, regPath, registry.QUERY_VALUE)
	if err != nil {
		return "", err
	}
	defer k.Close()
	v, _, err := k.GetStringValue("DeviceDesc")
	return v, err
}

// chooseGPU 让用户从列表中选择一块显卡
func chooseGPU(gpus []GPUInfo) (GPUInfo, error) {
	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Printf("请输入显卡序号 (1-%d)：", len(gpus))
		input, _ := reader.ReadString('\n')
		input = strings.TrimSpace(input)

		var idx int
		if _, err := fmt.Sscanf(input, "%d", &idx); err != nil {
			fmt.Println("输入无效，请输入数字。")
			continue
		}
		if idx < 1 || idx > len(gpus) {
			fmt.Printf("序号超出范围，请输入 1-%d。\n", len(gpus))
			continue
		}
		return gpus[idx-1], nil
	}
}

// readNewName 读取用户输入的新名称
func readNewName() string {
	reader := bufio.NewReader(os.Stdin)
	fmt.Print("请输入新的显卡名称（直接回车取消）：")
	input, _ := reader.ReadString('\n')
	return strings.TrimSpace(input)
}

// renameGPU 修改注册表中的 DeviceDesc
func renameGPU(g GPUInfo, newName string) error {
	k, err := registry.OpenKey(registry.LOCAL_MACHINE, g.Path, registry.SET_VALUE)
	if err != nil {
		return err
	}
	defer k.Close()

	return k.SetStringValue("DeviceDesc", newName)
}

// verifyGPU 重新读取 DeviceDesc 进行验证
func verifyGPU(g GPUInfo) (string, error) {
	desc, err := getDeviceDesc(g.Path)
	if err != nil {
		return "", err
	}
	return cleanDeviceDesc(desc), nil
}

// waitEnter 等待用户按键后退出
func waitEnter() {
	fmt.Println()
	fmt.Print("按回车键退出...")
	bufio.NewReader(os.Stdin).ReadString('\n')
}