package main

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"unsafe"
)

const fallbackDLLVersion = "4.1.33.6009"
const legacyLoaderSHA256 = "ce4557bf449fd53078f2eaa9becbf43b6f2bf40f16199e1a4bd6088ab233a65a"

type DLLBundle struct {
	HelperPath    string
	LoaderPath    string
	HelperVersion string
	LoaderVersion string
	WxWorkPath    string
	WxWorkVersion string
	Compatible    bool
	Message       string
}

type vsFixedFileInfo struct {
	Signature        uint32
	StrucVersion     uint32
	FileVersionMS    uint32
	FileVersionLS    uint32
	ProductVersionMS uint32
	ProductVersionLS uint32
	FileFlagsMask    uint32
	FileFlags        uint32
	FileOS           uint32
	FileType         uint32
	FileSubtype      uint32
	FileDateMS       uint32
	FileDateLS       uint32
}

func resolveDLLBundle() DLLBundle {
	exeDir := "."
	if exePath, err := os.Executable(); err == nil {
		exeDir = filepath.Dir(exePath)
	}

	wxWorkPath := detectWxWorkPath()
	if wxWorkPath == "" {
		wxWorkPath = `C:\企业微信\WXWork\WXWork.exe`
	}
	wxWorkVersion := ""
	if _, err := os.Stat(wxWorkPath); err == nil {
		wxWorkVersion, _ = getWindowsFileVersion(wxWorkPath)
	}

	versions := make([]string, 0, 4)
	if wxWorkVersion != "" {
		versions = append(versions, wxWorkVersion)
	}
	versions = append(versions, fallbackDLLVersion)
	versions = append(versions, scanDLLVersions(exeDir)...)

	seen := make(map[string]bool)
	for _, version := range versions {
		version = strings.TrimSpace(version)
		if version == "" || seen[version] {
			continue
		}
		seen[version] = true
		helperPath := filepath.Join(exeDir, "Helper_"+version+".dll")
		loaderPath := filepath.Join(exeDir, "Loader_"+version+".dll")
		if fileExists(helperPath) && fileExists(loaderPath) {
			compatible := wxWorkVersion != "" && sameVersionFamily(wxWorkVersion, version)
			message := ""
			if !compatible && wxWorkVersion != "" {
				message = fmt.Sprintf("WXWork %s is not compatible with Helper/Loader %s. Put Helper_%s.dll and Loader_%s.dll in build/bin to enable account/message callbacks.", wxWorkVersion, version, wxWorkVersion, wxWorkVersion)
			}
			if compatible && version != fallbackDLLVersion {
				fallbackLoaderPath := filepath.Join(exeDir, "Loader_"+fallbackDLLVersion+".dll")
				if sameFileContent(loaderPath, fallbackLoaderPath) || fileMatchesSHA256(loaderPath, legacyLoaderSHA256) {
					compatible = false
					message = fmt.Sprintf("Loader_%s.dll has the same content as Loader_%s.dll. Replace it with the real Loader_%s.dll before starting WXWork.", version, fallbackDLLVersion, version)
				}
			}
			return DLLBundle{
				HelperPath:    helperPath,
				LoaderPath:    loaderPath,
				HelperVersion: version,
				LoaderVersion: version,
				WxWorkPath:    wxWorkPath,
				WxWorkVersion: wxWorkVersion,
				Compatible:    compatible,
				Message:       message,
			}
		}
	}

	message := "No Helper/Loader DLL pair found."
	if wxWorkVersion != "" {
		message = fmt.Sprintf("No Helper_%s.dll and Loader_%s.dll found in %s.", wxWorkVersion, wxWorkVersion, exeDir)
	}
	return DLLBundle{
		WxWorkPath:    wxWorkPath,
		WxWorkVersion: wxWorkVersion,
		Message:       message,
	}
}

func detectWxWorkPath() string {
	if path := getRunningProcessImagePath("WXWork.exe"); path != "" {
		return path
	}
	if path, err := getWxWorkInstallPath(); err == nil && fileExists(path) {
		return path
	}
	candidates := []string{
		`C:\Program Files (x86)\WXWork\WXWork.exe`,
		`C:\Program Files\WXWork\WXWork.exe`,
		`C:\企业微信\WXWork\WXWork.exe`,
	}
	for _, path := range candidates {
		if fileExists(path) {
			return path
		}
	}
	return ""
}

func getRunningProcessImagePath(processName string) string {
	hSnapshot, err := syscall.CreateToolhelp32Snapshot(syscall.TH32CS_SNAPPROCESS, 0)
	if err != nil {
		return ""
	}
	defer syscall.CloseHandle(hSnapshot)

	var pe32 syscall.ProcessEntry32
	pe32.Size = uint32(unsafe.Sizeof(pe32))
	if err := syscall.Process32First(hSnapshot, &pe32); err != nil {
		return ""
	}

	for {
		name := syscall.UTF16ToString(pe32.ExeFile[:])
		if strings.EqualFold(name, processName) {
			if path := queryProcessImagePath(pe32.ProcessID); path != "" {
				return path
			}
		}
		if err := syscall.Process32Next(hSnapshot, &pe32); err != nil {
			break
		}
	}
	return ""
}

func queryProcessImagePath(pid uint32) string {
	const processQueryLimitedInformation = 0x1000
	handle, err := syscall.OpenProcess(processQueryLimitedInformation, false, pid)
	if err != nil {
		return ""
	}
	defer syscall.CloseHandle(handle)

	kernel32 := syscall.NewLazyDLL("kernel32.dll")
	queryFullProcessImageName := kernel32.NewProc("QueryFullProcessImageNameW")
	buf := make([]uint16, 32768)
	size := uint32(len(buf))
	ret, _, _ := queryFullProcessImageName.Call(
		uintptr(handle),
		0,
		uintptr(unsafe.Pointer(&buf[0])),
		uintptr(unsafe.Pointer(&size)),
	)
	if ret == 0 || size == 0 {
		return ""
	}
	return syscall.UTF16ToString(buf[:size])
}

func sameFileContent(pathA, pathB string) bool {
	infoA, errA := os.Stat(pathA)
	infoB, errB := os.Stat(pathB)
	if errA != nil || errB != nil || infoA.Size() != infoB.Size() {
		return false
	}
	hashA, errA := fileSHA256(pathA)
	hashB, errB := fileSHA256(pathB)
	return errA == nil && errB == nil && hashA == hashB
}

func fileMatchesSHA256(path string, expected string) bool {
	hash, err := fileSHA256(path)
	return err == nil && strings.EqualFold(hex.EncodeToString(hash[:]), expected)
}

func fileSHA256(path string) ([32]byte, error) {
	var zero [32]byte
	f, err := os.Open(path)
	if err != nil {
		return zero, err
	}
	defer f.Close()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return zero, err
	}
	var sum [32]byte
	copy(sum[:], h.Sum(nil))
	return sum, nil
}

func scanDLLVersions(dir string) []string {
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil
	}
	hasHelper := make(map[string]bool)
	hasLoader := make(map[string]bool)
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		name := entry.Name()
		lower := strings.ToLower(name)
		if strings.HasPrefix(lower, "helper_") && strings.HasSuffix(lower, ".dll") {
			hasHelper[versionFromDLLName(name)] = true
		}
		if strings.HasPrefix(lower, "loader_") && strings.HasSuffix(lower, ".dll") {
			hasLoader[versionFromDLLName(name)] = true
		}
	}
	versions := make([]string, 0)
	for version := range hasHelper {
		if version != "" && hasLoader[version] {
			versions = append(versions, version)
		}
	}
	return versions
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

func getWxWorkVersionDiagnostics() map[string]interface{} {
	bundle := resolveDLLBundle()
	return map[string]interface{}{
		"wxWorkPath":    bundle.WxWorkPath,
		"wxWorkVersion": bundle.WxWorkVersion,
		"helperPath":    bundle.HelperPath,
		"loaderPath":    bundle.LoaderPath,
		"helperVersion": bundle.HelperVersion,
		"loaderVersion": bundle.LoaderVersion,
		"compatible":    bundle.Compatible,
		"message":       bundle.Message,
	}
}

func versionFromDLLName(path string) string {
	name := filepath.Base(path)
	name = strings.TrimSuffix(name, filepath.Ext(name))
	idx := strings.LastIndex(name, "_")
	if idx < 0 || idx+1 >= len(name) {
		return ""
	}
	return strings.TrimSpace(name[idx+1:])
}

func sameVersionFamily(actual string, expected string) bool {
	actualParts := strings.Split(actual, ".")
	expectedParts := strings.Split(expected, ".")
	if len(actualParts) < 3 || len(expectedParts) < 3 {
		return actual == expected
	}
	return actualParts[0] == expectedParts[0] &&
		actualParts[1] == expectedParts[1] &&
		actualParts[2] == expectedParts[2]
}

func getWindowsFileVersion(path string) (string, error) {
	versionDLL := syscall.NewLazyDLL("version.dll")
	getFileVersionInfoSize := versionDLL.NewProc("GetFileVersionInfoSizeW")
	getFileVersionInfo := versionDLL.NewProc("GetFileVersionInfoW")
	verQueryValue := versionDLL.NewProc("VerQueryValueW")

	pathPtr, err := syscall.UTF16PtrFromString(path)
	if err != nil {
		return "", err
	}
	var handle uint32
	size, _, _ := getFileVersionInfoSize.Call(
		uintptr(unsafe.Pointer(pathPtr)),
		uintptr(unsafe.Pointer(&handle)),
	)
	if size == 0 {
		return "", fmt.Errorf("GetFileVersionInfoSizeW returned 0")
	}

	buf := make([]byte, size)
	ret, _, err := getFileVersionInfo.Call(
		uintptr(unsafe.Pointer(pathPtr)),
		0,
		size,
		uintptr(unsafe.Pointer(&buf[0])),
	)
	if ret == 0 {
		return "", fmt.Errorf("GetFileVersionInfoW failed: %v", err)
	}

	root, err := syscall.UTF16PtrFromString(`\`)
	if err != nil {
		return "", err
	}
	var block uintptr
	var blockLen uint32
	ret, _, err = verQueryValue.Call(
		uintptr(unsafe.Pointer(&buf[0])),
		uintptr(unsafe.Pointer(root)),
		uintptr(unsafe.Pointer(&block)),
		uintptr(unsafe.Pointer(&blockLen)),
	)
	if ret == 0 || block == 0 {
		return "", fmt.Errorf("VerQueryValueW failed: %v", err)
	}

	info := (*vsFixedFileInfo)(unsafe.Pointer(block))
	major := info.FileVersionMS >> 16
	minor := info.FileVersionMS & 0xffff
	build := info.FileVersionLS >> 16
	patch := info.FileVersionLS & 0xffff
	return fmt.Sprintf("%d.%d.%d.%d", major, minor, build, patch), nil
}