//go:build windows

package winpwn

import (
	"fmt"
	"strings"
	"sync"
	"unsafe"

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

// ListLoadedModules walks pid's PEB → Ldr → InMemoryOrderModuleList and
// returns the load address of every currently-loaded module, keyed by
// lower-cased base name ("kernel32.dll", "ntdll.dll", etc.).
//
// This is the single-call equivalent of calling ResolveModuleBase for every
// DLL in the process — use it when you need more than one or two bases, or
// when you want to enumerate what's loaded without knowing names in advance.
func ListLoadedModules(pid uint32) (map[string]uintptr, error) {
	h, err := windows.OpenProcess(
		windows.PROCESS_QUERY_INFORMATION|windows.PROCESS_VM_READ, false, pid)
	if err != nil {
		return nil, err
	}
	defer windows.CloseHandle(h)
	return ldrWalkAll(h, pid)
}

// ldrWalkAll performs the PEB→Ldr→InMemoryOrderModuleList walk and collects
// every entry, keyed by lower-cased base name. Shared by ListLoadedModules
// and ProcessSymbols.loadAll.
func ldrWalkAll(h windows.Handle, pid uint32) (map[string]uintptr, error) {
	var pbi windows.PROCESS_BASIC_INFORMATION
	var retLen uint32
	if err := windows.NtQueryInformationProcess(h, windows.ProcessBasicInformation,
		unsafe.Pointer(&pbi), uint32(unsafe.Sizeof(pbi)), &retLen); err != nil {
		return nil, fmt.Errorf("NtQueryInformationProcess: %w", err)
	}
	pebAddr := uintptr(unsafe.Pointer(pbi.PebBaseAddress))
	if pebAddr == 0 {
		return nil, fmt.Errorf("PEB is null for pid %d (not yet initialized?)", pid)
	}

	peb, err := readRemoteStruct[windows.PEB](h, pebAddr)
	if err != nil {
		return nil, fmt.Errorf("read PEB: %w", err)
	}
	ldrAddr := uintptr(unsafe.Pointer(peb.Ldr))
	if ldrAddr == 0 {
		return nil, fmt.Errorf("PEB.Ldr is null (loader not yet run in pid %d)", pid)
	}

	ldr, err := readRemoteStruct[windows.PEB_LDR_DATA](h, ldrAddr)
	if err != nil {
		return nil, fmt.Errorf("read PEB_LDR_DATA: %w", err)
	}

	entryLinksOffset := unsafe.Offsetof(windows.LDR_DATA_TABLE_ENTRY{}.InMemoryOrderLinks)
	headAddr := ldrAddr + unsafe.Offsetof(windows.PEB_LDR_DATA{}.InMemoryOrderModuleList)

	out := make(map[string]uintptr)
	cur := uintptr(unsafe.Pointer(ldr.InMemoryOrderModuleList.Flink))
	for cur != 0 && cur != headAddr {
		entryAddr := cur - entryLinksOffset
		entry, err := readRemoteStruct[windows.LDR_DATA_TABLE_ENTRY](h, entryAddr)
		if err != nil {
			return nil, fmt.Errorf("read LDR_DATA_TABLE_ENTRY: %w", err)
		}
		if name, err := readRemoteUTF16(h, entry.FullDllName); err == nil {
			key := strings.ToLower(moduleBaseName(name))
			if key != "" {
				out[key] = entry.DllBase
			}
		}
		cur = uintptr(unsafe.Pointer(entry.InMemoryOrderLinks.Flink))
	}
	return out, nil
}

// SymbolVA resolves the virtual address of a named export from a module
// currently loaded in pid. It combines ResolveModuleBase + OpenPEFromProcess
// + GetProcAddress into a single call — the Go equivalent of pwintools'
// p.symbols["kernel32.dll"]["WinExec"].
func SymbolVA(pid uint32, dll, name string) (uintptr, error) {
	base, err := ResolveModuleBase(pid, dll)
	if err != nil {
		return 0, err
	}
	pf, err := OpenPEFromProcess(pid, base)
	if err != nil {
		return 0, fmt.Errorf("open %s in pid %d: %w", dll, pid, err)
	}
	defer pf.Close()
	rva, err := pf.GetProcAddress(name)
	if err != nil {
		return 0, fmt.Errorf("%s!%s: %w", dll, name, err)
	}
	return base + uintptr(rva), nil
}

// ProcessSymbols resolves and caches loaded-module bases and exported symbol
// VAs for a running process — the Go equivalent of pwintools' p.libs /
// p.symbols. Caches one PEFile per DLL so repeated symbol lookups in the
// same module are cheap.
//
// Usage:
//
//	tube, _ := winpwn.Spawn("chal.exe")
//	sym := winpwn.NewProcessSymbols(tube.PID())
//	defer sym.Close()
//
//	k32, _       := sym.Base("kernel32.dll")
//	winexec, _   := sym.Symbol("kernel32.dll", "WinExec")
//	mods, _      := sym.Modules()   // all loaded DLLs
type ProcessSymbols struct {
	pid   uint32
	mu    sync.Mutex
	cache map[string]*symModule // keyed by lower-cased base name
}

type symModule struct {
	base uintptr
	pf   *PEFile
}

// NewProcessSymbols creates a ProcessSymbols for the given PID. No I/O
// happens until the first Base/Symbol call.
func NewProcessSymbols(pid uint32) *ProcessSymbols {
	return &ProcessSymbols{pid: pid, cache: make(map[string]*symModule)}
}

// Base returns the load address of the named module (e.g. "kernel32.dll"),
// the Go equivalent of pwintools' p.libs["kernel32.dll"].
func (ps *ProcessSymbols) Base(dll string) (uint64, error) {
	m, err := ps.loadModule(dll)
	if err != nil {
		return 0, err
	}
	return uint64(m.base), nil
}

// Symbol returns the virtual address of name exported from dll
// (e.g. "kernel32.dll", "WinExec"), the equivalent of pwintools'
// p.symbols["kernel32.dll"]["WinExec"].
func (ps *ProcessSymbols) Symbol(dll, name string) (uint64, error) {
	m, err := ps.loadModule(dll)
	if err != nil {
		return 0, err
	}
	rva, err := m.pf.GetProcAddress(name)
	if err != nil {
		return 0, fmt.Errorf("%s!%s: %w", dll, name, err)
	}
	return uint64(m.base) + uint64(rva), nil
}

// Modules returns a snapshot of every module currently loaded in the process,
// keyed by lower-cased base name — the equivalent of pwintools' p.libs dict.
func (ps *ProcessSymbols) Modules() (map[string]uint64, error) {
	raw, err := ListLoadedModules(ps.pid)
	if err != nil {
		return nil, err
	}
	out := make(map[string]uint64, len(raw))
	for k, v := range raw {
		out[k] = uint64(v)
	}
	return out, nil
}

// AllSymbols returns every named export from dll as a map of name → VA.
// Useful for quick "what's available" exploration without knowing exact names.
func (ps *ProcessSymbols) AllSymbols(dll string) (map[string]uint64, error) {
	m, err := ps.loadModule(dll)
	if err != nil {
		return nil, err
	}
	exports, err := m.pf.ListExports()
	if err != nil {
		return nil, err
	}
	out := make(map[string]uint64, len(exports))
	for _, e := range exports {
		if e.Name != "" && e.RVA != 0 {
			out[e.Name] = uint64(m.base) + uint64(e.RVA)
		}
	}
	return out, nil
}

// Close releases all cached PEFile handles. Safe to call more than once.
func (ps *ProcessSymbols) Close() {
	ps.mu.Lock()
	defer ps.mu.Unlock()
	for _, m := range ps.cache {
		if m.pf != nil {
			m.pf.Close()
		}
	}
	ps.cache = nil
}

func (ps *ProcessSymbols) loadModule(dll string) (*symModule, error) {
	key := strings.ToLower(moduleBaseName(dll))
	ps.mu.Lock()
	defer ps.mu.Unlock()
	if ps.cache == nil {
		return nil, fmt.Errorf("ProcessSymbols already closed")
	}
	if m, ok := ps.cache[key]; ok {
		return m, nil
	}
	base, err := ResolveModuleBase(ps.pid, dll)
	if err != nil {
		return nil, err
	}
	pf, err := OpenPEFromProcess(ps.pid, base)
	if err != nil {
		return nil, fmt.Errorf("open %s in pid %d: %w", dll, ps.pid, err)
	}
	m := &symModule{base: base, pf: pf}
	ps.cache[key] = m
	return m, nil
}