//go:build windows
// +build windows

package iptables

import (
	"fmt"
	"os/exec"
	"sync"
	"syscall"
	"time"
	"unsafe"

	"github.com/v2rayA/v2rayA/pkg/util/log"
	"golang.org/x/sys/windows"
	"golang.org/x/sys/windows/registry"
)

const (
	errnoERROR_IO_PENDING            = 997
	INTERNET_OPTION_SETTINGS_CHANGED = 39
)

var (
	errERROR_IO_PENDING    error = syscall.Errno(errnoERROR_IO_PENDING)
	modwininet                   = syscall.NewLazyDLL("wininet.dll")
	procInternetSetOptionW       = modwininet.NewProc("InternetSetOptionW")
)

// errnoErr returns common boxed Errno values, to prevent
// allocations at runtime.
func errnoErr(e syscall.Errno) error {
	switch e {
	case 0:
		return nil
	case errnoERROR_IO_PENDING:
		return errERROR_IO_PENDING
	}
	// TODO: add more here, after collecting data on the common
	// error values see on Windows. (perhaps when running
	// all.bat?)
	return e
}

var (
	modadvapi32 = windows.NewLazySystemDLL("advapi32.dll")
	modkernel32 = windows.NewLazySystemDLL("kernel32.dll")

	procGetCurrentThread = modkernel32.NewProc("GetCurrentThread")
	procOpenThreadToken  = modadvapi32.NewProc("OpenThreadToken")
	procImpersonateSelf  = modadvapi32.NewProc("ImpersonateSelf")
	procRevertToSelf     = modadvapi32.NewProc("RevertToSelf")
)

func GetCurrentThread() (pseudoHandle windows.Handle, err error) {
	r0, _, e1 := syscall.SyscallN(procGetCurrentThread.Addr())
	pseudoHandle = windows.Handle(r0)
	if pseudoHandle == 0 {
		if e1 != 0 {
			err = errnoErr(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

func OpenThreadToken(h windows.Handle, access uint32, self bool, token *windows.Token) (err error) {
	var _p0 uint32
	if self {
		_p0 = 1
	} else {
		_p0 = 0
	}
	r1, _, e1 := syscall.SyscallN(procOpenThreadToken.Addr(), uintptr(h), uintptr(access), uintptr(_p0), uintptr(unsafe.Pointer(token)))
	if r1 == 0 {
		if e1 != 0 {
			err = errnoErr(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

func ImpersonateSelf() (err error) {
	r0, _, e1 := syscall.SyscallN(procImpersonateSelf.Addr(), uintptr(2))
	if r0 == 0 {
		if e1 != 0 {
			err = errnoErr(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

func RevertToSelf() (err error) {
	r0, _, e1 := syscall.SyscallN(procRevertToSelf.Addr())
	if r0 == 0 {
		if e1 != 0 {
			err = errnoErr(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

func OpenCurrentThreadToken() (windows.Token, error) {
	if e := ImpersonateSelf(); e != nil {
		return 0, e
	}
	defer RevertToSelf()
	t, e := GetCurrentThread()
	if e != nil {
		return 0, e
	}
	var tok windows.Token
	e = OpenThreadToken(t, windows.TOKEN_QUERY, true, &tok)
	if e != nil {
		return 0, e
	}
	return tok, nil
}

func HasAdminRights() (bool, error) {
	// https://github.com/golang/go/issues/28804
	var sid *windows.SID
	err := windows.AllocateAndInitializeSid(&windows.SECURITY_NT_AUTHORITY, 2, windows.SECURITY_BUILTIN_DOMAIN_RID, windows.DOMAIN_ALIAS_RID_ADMINS, 0, 0, 0, 0, 0, 0, &sid)
	if err != nil {
		return false, err
	}
	defer windows.FreeSid(sid)

	token, err := OpenCurrentThreadToken()
	if err != nil {
		return false, err
	}

	member, err := token.IsMember(sid)
	if err != nil {
		return false, err
	}
	return member, nil
}

func InternetOptionSettingsChanged() (syscall.Handle, error) {
	// https://github.com/Diving-Fish/maimaidx-prober/blob/1673cff975e6bbb14eeeda6260af131ac036f974/proxy/lib/wininet_windows.go#L18
	// https://docs.microsoft.com/en-us/windows/win32/wininet/option-flags
	p1 := uint16(0)
	p2 := uint64(INTERNET_OPTION_SETTINGS_CHANGED)
	p3 := uint16(0)
	p4 := uint64(0)
	r1, _, e1 := syscall.SyscallN(
		procInternetSetOptionW.Addr(),
		uintptr(p1),
		uintptr(p2),
		uintptr(p3),
		uintptr(p4),
	)
	if r1 == 0 {
		if e1 != 0 {
			return 0, e1
		} else {
			return 0, syscall.EINVAL
		}
	}
	return syscall.Handle(r1), nil
}

const profileListPath = "SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\ProfileList"

func getProfileListSubKeyNames() ([]string, error) {
	// https://github.com/ninoseki/ninoseki.github.io/blob/0e911fe146c0de7b2fa4449ff9c2ac1ebe755b67/docs/_posts/2020-07-25-windows-registry-with-go.md
	k, err := registry.OpenKey(registry.LOCAL_MACHINE, profileListPath, registry.ENUMERATE_SUB_KEYS)
	if err != nil {
		return nil, fmt.Errorf("unable to open registry key %q: %v", profileListPath, err)
	}
	defer k.Close()
	return k.ReadSubKeyNames(0)
}

// windowsProxySavedState stores the original proxy values before v2rayA modifies them
type windowsProxySavedState struct {
	mu          sync.Mutex
	saved       bool
	proxyEnable uint32
	proxyServer string
}

var savedWindowsProxy windowsProxySavedState

type systemProxy struct{}

var SystemProxy systemProxy

func (p *systemProxy) AddIPWhitelist(cidr string) {}

func (p *systemProxy) RemoveIPWhitelist(cidr string) {}

type todo struct {
	Key    registry.Key
	Prefix string
}

// readRegistryProxyState reads ProxyEnable and ProxyServer from a registry key path
func readRegistryProxyState(key registry.Key, prefix string) (proxyEnable uint32, proxyServer string, err error) {
	k, err := registry.OpenKey(key, prefix+`SOFTWARE\Microsoft\Windows\CurrentVersion\Internet Settings`, registry.READ)
	if err != nil {
		return 0, "", fmt.Errorf("open key for reading: %v", err)
	}
	defer k.Close()

	val, _, err := k.GetIntegerValue("ProxyEnable")
	if err != nil {
		proxyEnable = 0
	} else {
		proxyEnable = uint32(val)
	}

	proxyServer, _, err = k.GetStringValue("ProxyServer")
	if err != nil {
		proxyServer = ""
	}

	return proxyEnable, proxyServer, nil
}

// saveProxyState reads and saves the current proxy state before we modify it
func saveProxyState(hasAdminRights bool) {
	var savedEnable uint32
	var savedServer string

	if hasAdminRights {
		sids, err := getProfileListSubKeyNames()
		if err != nil {
			log.Debug("saveProxyState: getProfileListSubKeyNames: %v", err)
			return
		}
		for _, sid := range sids {
			enable, server, err := readRegistryProxyState(registry.USERS, sid+`\`)
			if err != nil {
				log.Debug("saveProxyState: readRegistryProxyState for %v: %v", sid, err)
				continue
			}
			savedEnable = enable
			savedServer = server
			break
		}
	} else {
		enable, server, err := readRegistryProxyState(registry.CURRENT_USER, "")
		if err != nil {
			log.Debug("saveProxyState: readRegistryProxyState: %v", err)
			return
		}
		savedEnable = enable
		savedServer = server
	}

	savedWindowsProxy.mu.Lock()
	savedWindowsProxy.proxyEnable = savedEnable
	savedWindowsProxy.proxyServer = savedServer
	savedWindowsProxy.saved = true
	savedWindowsProxy.mu.Unlock()
}

func (p *systemProxy) GetSetupCommands() Setter {
	hasAdminRights, err := HasAdminRights()
	if err != nil {
		return NewErrorSetter(err)
	}
	setter := Setter{
		PreFunc: func() error {
			// Step 1: Save current proxy state before modifying
			saveProxyState(hasAdminRights)

			// Step 2: Original setup logic - set proxy to v2rayA
			var todolist []todo
			if hasAdminRights {
				// https://docs.microsoft.com/en-us/windows/win32/services/services-and-the-registry
				// A service should not access HKEY_CURRENT_USER or HKEY_CLASSES_ROOT, especially when impersonating a user.
				sids, err := getProfileListSubKeyNames()
				if err != nil {
					log.Debug("GetSetupCommands: getProfileListSubKeyNames: %v", err)
					return err
				}
				for _, sid := range sids {
					todolist = append(todolist, todo{
						Key:    registry.USERS,
						Prefix: sid + `\`,
					})
				}
			} else {
				todolist = append(todolist, todo{
					Key:    registry.CURRENT_USER,
					Prefix: "",
				})
			}
			for _, todo := range todolist {
				key, err := registry.OpenKey(todo.Key, todo.Prefix+`SOFTWARE\Microsoft\Windows\CurrentVersion\Internet Settings`, registry.ALL_ACCESS)
				if err != nil {
					log.Debug("GetSetupCommands: OpenKey %v, %v: %v", todo.Key, todo.Prefix+`SOFTWARE\Microsoft\Windows\CurrentVersion\Internet Settings`, err)
					continue
				}
				defer key.Close()
				_ = key.DeleteValue("AutoConfigURL")
				if err = key.SetDWordValue("ProxyEnable", 1); err != nil {
					log.Debug("GetSetupCommands: key: %v: SetDWordValue ProxyEnable: %v", todo.Key, err)
					return err
				}
				if err = key.SetStringValue("ProxyServer", "127.0.0.1:52345"); err != nil {
					log.Debug("GetSetupCommands: key: %v: SetStringValue ProxyServer: %v", todo.Key, err)
					return err
				}
			}
			_, _ = InternetOptionSettingsChanged()
			if hasAdminRights {
				// https://helpcenter.gsx.com/hc/en-us/articles/216487418-How-to-Import-Internet-Explorer-Proxy-Configuration-for-PowerShell-Use
				// You can browse the Internet and open OWA successfully using Internet Explorer (IE) but you cannot connect to Office 365 using PowerShell.
				// To fix this, we set Windows Proxy settings using NETSH for all applications that rely on default system configuration.
				time.Sleep(time.Second)
				return exec.Command("netsh", "winhttp", "import", "proxy", "source=ie").Run()
			} else {
				return nil
			}
		},
	}
	return setter
}

func (p *systemProxy) GetCleanCommands() Setter {
	hasAdminRights, err := HasAdminRights()
	if err != nil {
		return NewErrorSetter(err)
	}

	savedWindowsProxy.mu.Lock()
	saved := savedWindowsProxy.saved
	savedEnable := savedWindowsProxy.proxyEnable
	savedServer := savedWindowsProxy.proxyServer
	savedWindowsProxy.mu.Unlock()

	if !saved {
		// No saved state: fall back to original behavior (disable proxy)
		setter := Setter{
			PreFunc: func() error {
				var todolist []todo
				if hasAdminRights {
					sids, err := getProfileListSubKeyNames()
					if err != nil {
						return err
					}
					for _, sid := range sids {
						todolist = append(todolist, todo{
							Key:    registry.USERS,
							Prefix: sid + `\`,
						})
					}
				} else {
					todolist = append(todolist, todo{
						Key:    registry.CURRENT_USER,
						Prefix: "",
					})
				}

				var errs []error
				for _, todo := range todolist {
					key, _, err := registry.CreateKey(todo.Key, todo.Prefix+`SOFTWARE\Microsoft\Windows\CurrentVersion\Internet Settings`, registry.ALL_ACCESS)
					if err != nil {
						errs = append(errs, err)
						continue
					}
					defer key.Close()
					if err = key.SetDWordValue("ProxyEnable", 0); err != nil {
						errs = append(errs, err)
						continue
					}
				}
				_, _ = InternetOptionSettingsChanged()
				if hasAdminRights {
					e := exec.Command("netsh", "winhttp", "import", "proxy", "source=ie").Run()
					if e != nil {
						errs = append(errs, e)
					}
				}
				if len(errs) > 0 {
					return errs[0]
				}
				return nil
			},
		}
		return setter
	}

	// Restore original state
	setter := Setter{
		PreFunc: func() error {
			var todolist []todo
			if hasAdminRights {
				sids, err := getProfileListSubKeyNames()
				if err != nil {
					return err
				}
				for _, sid := range sids {
					todolist = append(todolist, todo{
						Key:    registry.USERS,
						Prefix: sid + `\`,
					})
				}
			} else {
				todolist = append(todolist, todo{
					Key:    registry.CURRENT_USER,
					Prefix: "",
				})
			}

			var errs []error
			for _, todo := range todolist {
				key, _, err := registry.CreateKey(todo.Key, todo.Prefix+`SOFTWARE\Microsoft\Windows\CurrentVersion\Internet Settings`, registry.ALL_ACCESS)
				if err != nil {
					errs = append(errs, err)
					continue
				}
				defer key.Close()

				if savedEnable == 1 {
					// Proxy was originally enabled: restore original server
					if savedServer != "" {
						if err = key.SetStringValue("ProxyServer", savedServer); err != nil {
							errs = append(errs, err)
							continue
						}
					}
					if err = key.SetDWordValue("ProxyEnable", 1); err != nil {
						errs = append(errs, err)
						continue
					}
				} else {
					// Proxy was originally disabled: disable it
					if err = key.SetDWordValue("ProxyEnable", 0); err != nil {
						errs = append(errs, err)
						continue
					}
				}
			}
			_, _ = InternetOptionSettingsChanged()
			if hasAdminRights {
				e := exec.Command("netsh", "winhttp", "import", "proxy", "source=ie").Run()
				if e != nil {
					errs = append(errs, e)
				}
			}
			if len(errs) > 0 {
				return errs[0]
			}
			return nil
		},
	}

	savedWindowsProxy.mu.Lock()
	savedWindowsProxy.saved = false
	savedWindowsProxy.mu.Unlock()

	return setter
}
