package main

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"time"
	"unsafe"

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

// 进程相关的Windows API常量和结构
const (
	TH32CS_SNAPPROCESS = 0x00000002
)

// getUserWxWorkVersion 获取企业微信版本
func getUserWxWorkVersion(funcAddr uintptr) (string, error) {
	// 首先检查函数指针是否有效
	if funcAddr == 0 {
		return "", fmt.Errorf("GetUserWxWorkVersion函数指针无效")
	}

	// 调用GetUserWxWorkVersion函数前添加defer recover，防止程序崩溃
	done := false
	defer func() {
		if r := recover(); r != nil {
			if !done {
				globalLogger.Error("调用GetUserWxWorkVersion时发生panic: %v", r)
			}
		}
	}()

	// 调用GetUserWxWorkVersion函数
	versionPtr, _, callErr := syscall.Syscall(
		funcAddr,
		0,
		0,
		0,
		0,
	)

	if versionPtr == 0 || callErr != 0 {
		done = true
		return "", fmt.Errorf("调用GetUserWxWorkVersion失败: 返回值=%d, 错误=%v", versionPtr, callErr)
	}

	// 安全地将指针转换为Go字符串
	version := ""
	if versionPtr != 0 {
		// 使用defer recover保护指针转换操作
		defer func() {
			if r := recover(); r != nil {
				done = true
				globalLogger.Error("转换版本字符串时发生panic: %v", r)
				version = ""
			}
		}()
		version = syscall.UTF16ToString((*[1024]uint16)(unsafe.Pointer(versionPtr))[:])
	}
	done = true
	return version, nil
}

// findProcessByName 根据进程名查找进程ID
func findProcessByName(name string) ([]uint32, error) {
	globalLogger.Info("[辅助程序] 查找进程: %s", name)

	// 打开系统快照
	hSnapshot, err := syscall.CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)
	if err != nil {
		return nil, fmt.Errorf("创建进程快照失败: %v", err)
	}
	defer syscall.CloseHandle(hSnapshot)

	// 初始化进程信息结构 - 使用syscall包中的ProcessEntry32
	var pe32 syscall.ProcessEntry32
	pe32.Size = uint32(unsafe.Sizeof(pe32))

	// 获取第一个进程
	if err := syscall.Process32First(hSnapshot, &pe32); err != nil {
		return nil, fmt.Errorf("获取第一个进程失败: %v", err)
	}

	// 存储找到的进程ID
	var processIDs []uint32

	// 遍历所有进程
	for {
		// 将进程名从UTF-16转换为Go字符串
		processName := syscall.UTF16ToString(pe32.ExeFile[:])

		// 只匹配完整的进程名，不区分大小写
		if strings.EqualFold(processName, name) {

			globalLogger.Info("[辅助程序] 找到进程: %s, 进程ID: %d", processName, pe32.ProcessID)
			processIDs = append(processIDs, pe32.ProcessID)
		}

		// 获取下一个进程
		if err := syscall.Process32Next(hSnapshot, &pe32); err != nil {
			// 如果没有更多进程，跳出循环
			if err == syscall.ERROR_NO_MORE_FILES {
				break
			}
			return nil, fmt.Errorf("获取下一个进程失败: %v", err)
		}
	}

	if len(processIDs) == 0 {
		return []uint32{}, fmt.Errorf("未找到进程: %s", name)
	}

	return processIDs, nil
}

// getWxWorkInstallPath 从注册表获取企业微信安装路径
func getWxWorkInstallPath() (string, error) {
	// 打开注册表项
	key, err := registry.OpenKey(registry.LOCAL_MACHINE,
		"SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\WXWork",
		registry.QUERY_VALUE)
	if err != nil {
		return "", fmt.Errorf("无法打开注册表: %v", err)
	}
	defer key.Close()

	// 读取安装路径
	installPath, _, err := key.GetStringValue("InstallLocation")
	if err != nil {
		return "", fmt.Errorf("无法获取安装路径: %v", err)
	}

	// 验证路径是否存在
	exePath := filepath.Join(installPath, "WXWork.exe")
	if _, err := os.Stat(exePath); os.IsNotExist(err) {
		return "", fmt.Errorf("企业微信可执行文件不存在: %s", exePath)
	}

	return exePath, nil
}

// startWxWorkInstance 启动企业微信新实例
func startWxWorkInstance(exePath string) (uint32, error) {
	// 创建进程信息
	var si syscall.StartupInfo
	var pi syscall.ProcessInformation
	defer syscall.CloseHandle(pi.Thread)
	defer syscall.CloseHandle(pi.Process)

	// 准备命令行参数
	cmdLine, err := syscall.UTF16PtrFromString(exePath + " --multi-instance")
	if err != nil {
		return 0, fmt.Errorf("转换命令行参数失败: %v", err)
	}

	// 启动进程
	err = syscall.CreateProcess(nil, cmdLine, nil, nil, false,
		syscall.CREATE_NEW_PROCESS_GROUP, nil, nil, &si, &pi)
	if err != nil {
		return 0, fmt.Errorf("启动进程失败: %v", err)
	}

	globalLogger.Info("成功启动企业微信实例，进程ID: %d", pi.ProcessId)
	return pi.ProcessId, nil
}

// InjectWxWork 智能多开并注入DLL
// 参数1: WxWorkHelper.dll路径
// 参数2: WXWork.exe路径，传空
// 返回值: 成功返回企业微信进程ID，失败返回0
// 函数原型: DWORD __stdcall InjectWxWork(IN LPCSTR szDllPath，IN LPCSTR szWxWorkExePath);
func InjectWxWork(szDllPath, szWxWorkExePath string) uint32 {
	// 1. 检查DLL文件是否存在
	if _, err := os.Stat(szDllPath); os.IsNotExist(err) {
		globalLogger.Error("DLL文件不存在: %s", szDllPath)
		return 0
	}

	// 2. 加载Loader DLL获取注入函数
	dllBundle := resolveDLLBundle()
	loaderPath := dllBundle.LoaderPath
	if loaderPath == "" {
		globalLogger.Error("未找到可用Loader DLL: %s", dllBundle.Message)
		return 0
	}
	loaderDLL, err := loadDLL(loaderPath)
	if err != nil {
		globalLogger.Error("加载Loader DLL失败: %v", err)
		return 0
	}
	defer syscall.FreeLibrary(loaderDLL)

	// 3. 获取Loader函数指针
	loaderFuncs, err := getLoaderFunctions(loaderDLL)
	if err != nil || loaderFuncs.InjectWxWork == 0 {
		globalLogger.Error("获取InjectWxWork函数指针失败: %v", err)
		return 0
	}

	// 4. 将参数转换为ANSI字符串指针(LPCSTR)，以匹配InjectWxWork函数的参数要求
	dllPathPtr, err := syscall.BytePtrFromString(szDllPath)
	if err != nil {
		globalLogger.Error("转换DLL路径失败: %v", err)
		return 0
	}

	// 企业微信可执行文件路径参数，传空时使用空指针
	var wxWorkExePathPtr *byte
	if szWxWorkExePath != "" {
		wxWorkExePathPtr, err = syscall.BytePtrFromString(szWxWorkExePath)
		if err != nil {
			globalLogger.Error("转换企业微信路径失败: %v", err)
			return 0
		}
	}

	// 5. 直接使用syscall.Syscall调用InjectWxWork函数进行智能多开并注入DLL
	globalLogger.Info("准备调用InjectWxWork函数，DLL路径: %s，企业微信路径: %s，函数地址: 0x%X",
		szDllPath, szWxWorkExePath, loaderFuncs.InjectWxWork)
	ret, _, callErr := syscall.Syscall(
		loaderFuncs.InjectWxWork,            // 函数地址
		2,                                   // 参数数量
		uintptr(unsafe.Pointer(dllPathPtr)), // 参数1: DLL路径 (WxWorkHelper.dll)
		uintptr(unsafe.Pointer(wxWorkExePathPtr)), // 参数2: 企业微信路径 (传空时为NULL)
		0, // 保留参数
	)

	// 6. 检查函数调用结果
	if ret == 0 {
		// 调用失败，记录详细错误信息
		globalLogger.Error("调用InjectWxWork函数失败，返回值: %d, 错误代码: %v", ret, callErr)
		return 0
	}

	globalLogger.Info("InjectWxWork函数调用成功，返回进程ID: %d", ret)
	return uint32(ret)
}

func startAdditionalWxWorkInstance(helperDLLPath string) (map[string]interface{}, error) {
	if strings.TrimSpace(helperDLLPath) == "" {
		return nil, fmt.Errorf("helper DLL path is empty")
	}
	if _, err := os.Stat(helperDLLPath); err != nil {
		return nil, fmt.Errorf("helper DLL is unavailable: %w", err)
	}

	loaderFuncsMutex.Lock()
	loaderFuncs := globalLoaderFuncs
	loaderFuncsMutex.Unlock()
	if loaderFuncs == nil {
		return nil, fmt.Errorf("loader functions are not initialized")
	}

	if loaderFuncs.InjectWxWorkMultiOpen != 0 {
		if pid, err := injectWxWorkMultiOpen(loaderFuncs.InjectWxWorkMultiOpen, helperDLLPath, ""); err == nil && pid != 0 {
			setInjectionStatus(injectionStatusConnected, pid, "")
			go injectAllWxWorkProcesses(helperDLLPath, pid)
			return map[string]interface{}{
				"success":               true,
				"message":               "new WeCom instance requested by multi-open",
				"method":                "InjectWxWorkMultiOpen",
				"processId":             pid,
				"recognizedClientCount": recognizedClientCount(),
				"usableClientCount":     usableClientCount(),
				"connectionCount":       connectedClientCount(),
			}, nil
		} else if err != nil {
			globalLogger.Warn("[辅助程序] InjectWxWorkMultiOpen failed, fallback to process start: %v", err)
		}
	}

	exePath, err := getWxWorkInstallPath()
	if err != nil {
		return nil, err
	}
	pid, err := startWxWorkInstance(exePath)
	if err != nil {
		return nil, err
	}
	if loaderFuncs.InjectWxWorkPid != 0 {
		if ok, injectErr := injectIntoProcess(loaderFuncs.InjectWxWorkPid, pid, helperDLLPath); injectErr != nil {
			globalLogger.Warn("[辅助程序] fallback instance started but InjectWxWorkPid failed: %v", injectErr)
		} else if ok {
			markPIDInjected(pid)
		}
	}
	go injectAllWxWorkProcesses(helperDLLPath, pid)
	setInjectionStatus(injectionStatusConnected, pid, "")
	return map[string]interface{}{
		"success":               true,
		"message":               "new WeCom instance started by fallback",
		"method":                "CreateProcess+InjectWxWorkPid",
		"processId":             pid,
		"recognizedClientCount": recognizedClientCount(),
		"usableClientCount":     usableClientCount(),
		"connectionCount":       connectedClientCount(),
	}, nil
}

func injectWxWorkMultiOpen(funcAddr uintptr, dllPath string, wxWorkExePath string) (uint32, error) {
	dllPathPtr, err := syscall.BytePtrFromString(dllPath)
	if err != nil {
		return 0, err
	}
	var wxWorkExePathPtr *byte
	if strings.TrimSpace(wxWorkExePath) != "" {
		wxWorkExePathPtr, err = syscall.BytePtrFromString(wxWorkExePath)
		if err != nil {
			return 0, err
		}
	}
	ret, _, callErr := syscall.Syscall(
		funcAddr,
		2,
		uintptr(unsafe.Pointer(dllPathPtr)),
		uintptr(unsafe.Pointer(wxWorkExePathPtr)),
		0,
	)
	if ret == 0 {
		return 0, fmt.Errorf("InjectWxWorkMultiOpen returned 0: %v", callErr)
	}
	return uint32(ret), nil
}

// injectAllWxWorkProcesses tries to inject the helper DLL into every WXWork.exe
// process. Enterprise WeChat often keeps several WXWork.exe processes alive; the
// PID returned by InjectWxWork is not always the one that owns message callbacks.
func injectAllWxWorkProcesses(dllPath string, primaryPID uint32) {
	time.Sleep(1500 * time.Millisecond)

	loaderFuncsMutex.Lock()
	loaderFuncs := globalLoaderFuncs
	loaderFuncsMutex.Unlock()
	if loaderFuncs == nil || loaderFuncs.InjectWxWorkPid == 0 {
		globalLogger.Warn("[辅助程序] 跳过多进程注入，InjectWxWorkPid函数不可用")
		return
	}

	processIDs, err := findProcessByName("WXWork.exe")
	if err != nil {
		globalLogger.Warn("[辅助程序] 多进程注入时未找到WXWork.exe: %v", err)
		return
	}
	globalLogger.Info("[辅助程序] 准备对 %d 个WXWork.exe进程尝试注入，主PID: %d", len(processIDs), primaryPID)

	successCount := 0
	for _, pid := range processIDs {
		if !markPIDInjected(pid) {
			globalLogger.Info("[辅助程序] WXWork进程 %d 已注入过，跳过", pid)
			continue
		}
		ok, injectErr := injectIntoProcess(loaderFuncs.InjectWxWorkPid, pid, dllPath)
		if injectErr != nil {
			unmarkPIDInjected(pid)
			globalLogger.Warn("[辅助程序] WXWork进程 %d 注入失败: %v", pid, injectErr)
			continue
		}
		if ok {
			successCount++
		}
	}
	globalLogger.Info("[辅助程序] WXWork多进程注入完成，成功: %d/%d", successCount, len(processIDs))
}

// injectIntoProcess 向指定进程注入DLL
// injectIntoProcess 调用InjectWxWorkPid函数将DLL注入到指定的企业微信进程
// 参数:
// - funcAddr: InjectWxWorkPid函数地址
// - pid: 目标企业微信进程ID
// - dllPath: DLL文件路径
// 返回值:
// - 成功与否的布尔值
// - 错误信息
func injectIntoProcess(funcAddr uintptr, pid uint32, dllPath string) (bool, error) {
	fmt.Printf("找到要注入的DLL文件: %s\n", dllPath)
	fmt.Printf("目标进程ID: %d\n", pid)

	// 验证DLL文件是否存在
	if _, err := os.Stat(dllPath); os.IsNotExist(err) {
		globalLogger.Error("[辅助程序] DLL文件不存在: %s", dllPath)
		return false, fmt.Errorf("DLL文件不存在: %s", dllPath)
	}

	// 检查函数地址是否有效
	if funcAddr == 0 {
		globalLogger.Error("[辅助程序] InjectWxWorkPid函数地址无效")
		return false, fmt.Errorf("InjectWxWorkPid函数地址无效")
	}

	// 验证目标进程是否存在
	systemProcess, err := os.FindProcess(int(pid))
	if err != nil {
		globalLogger.Error("[辅助程序] 找不到进程ID %d: %v", pid, err)
		return false, fmt.Errorf("找不到目标进程: %v", err)
	}
	systemProcess.Release() // 只是检查进程是否存在，不保留句柄

	// 将DLL路径转换为ANSI字符串指针(LPCSTR)，以匹配InjectWxWorkPid函数的参数要求
	dllPathPtr, err := syscall.BytePtrFromString(dllPath)
	if err != nil {
		fmt.Printf("错误: 无法转换DLL路径: %v\n", err)
		globalLogger.Error("[辅助程序] 转换DLL路径失败: %v", err)
		return false, fmt.Errorf("转换DLL路径失败: %v", err)
	}

	globalLogger.Info("[辅助程序] 准备注入DLL，路径: %s，目标进程ID: %d，函数地址: 0x%X", dllPath, pid, funcAddr)
	fmt.Printf("函数地址: 0x%X, 参数1(pid): %d, 参数2(dllPath): %s\n", funcAddr, pid, dllPath)

	// Windows API权限常量
	const (
		PROCESS_QUERY_INFORMATION = 0x0400
		PROCESS_VM_OPERATION      = 0x0008
		PROCESS_VM_READ           = 0x0010
		PROCESS_VM_WRITE          = 0x0020
		// Windows错误代码常量
		ERROR_ACCESS_DENIED     = 5
		ERROR_INVALID_PARAMETER = 87
		ERROR_NOT_FOUND         = 1168
	)

	// 检查DLL文件大小
	dllInfo, err := os.Stat(dllPath)
	if err == nil {
		globalLogger.Info("[辅助程序] DLL文件大小: %d 字节", dllInfo.Size())
	}

	// 调用InjectWxWorkPid函数(DWORD __stdcall InjectWxWorkPid(IN DWORD dwPid, IN LPCSTR szDllPath))
	// 使用与成功示例完全一致的syscall.Syscall调用方式
	fmt.Println("开始调用函数...")
	globalLogger.Info("[辅助程序] 开始调用InjectWxWorkPid函数，函数地址: 0x%X, 进程ID: %d, DLL路径: %s", funcAddr, pid, dllPath)
	ret, _, callErr := syscall.Syscall(
		funcAddr,                            // 函数地址
		2,                                   // 参数数量
		uintptr(pid),                        // 参数1: 进程ID
		uintptr(unsafe.Pointer(dllPathPtr)), // 参数2: DLL路径
		0,                                   // 保留参数
	)

	// 检查返回值和错误信息
	if ret == 0 {
		// 注入失败，记录详细错误信息
		fmt.Printf("注入DLL失败，返回值: %d, 错误代码: %v\n", ret, callErr)
		globalLogger.Error("[辅助程序] 注入DLL失败，返回值: %d, 错误代码: %v", ret, callErr)

		// 提供详细的错误分析和可能的解决方案
		possibleSolutions := "可能的解决方案:\n"
		possibleSolutions += "1. 确认以管理员权限运行程序\n"
		possibleSolutions += "2. 检查目标进程是否正在运行且未被保护\n"
		possibleSolutions += "3. 确保DLL文件未被占用且路径正确\n"
		possibleSolutions += "4. 确认DLL与目标进程的位数匹配(32位/64位)\n"
		possibleSolutions += "5. 检查企业微信版本是否受支持(3.0.0.1001及以上)"
		fmt.Println(possibleSolutions)
		globalLogger.Info("%s", possibleSolutions)

		return false, fmt.Errorf("注入DLL失败: 返回值=%d, 错误=%v\n%s", ret, callErr, possibleSolutions)
	} else {
		// 注入成功
		fmt.Printf("向进程ID %d 注入DLL成功，返回值: %d\n", pid, ret)
		globalLogger.Info("[辅助程序] 向进程ID %d 注入DLL成功，返回值: %d", pid, ret)
		return true, nil
	}
}