package winpwn

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"os"
	"unicode/utf16"
)

// minidumpSignature is MINIDUMP_HEADER.Signature ('MDMP' read as a
// little-endian ULONG32), the magic number every .dmp file starts with.
const minidumpSignature = 0x504D444D

// MinidumpStreamType mirrors winnt.h's MINIDUMP_STREAM_TYPE. Only the
// values this package decodes natively are named here; RawStream accepts
// any numeric stream type for everything else (SystemInfoStream,
// ThreadListStream, Memory64ListStream, ...), the same way dbghelp's
// MiniDumpReadDumpStream takes an arbitrary stream number.
type MinidumpStreamType uint32

const (
	StreamThreadList   MinidumpStreamType = 3
	StreamModuleList   MinidumpStreamType = 4
	StreamMemoryList   MinidumpStreamType = 5
	StreamException    MinidumpStreamType = 6
	StreamSystemInfo   MinidumpStreamType = 7
	StreamMemory64List MinidumpStreamType = 9
)

// minidumpHeader mirrors MINIDUMP_HEADER (winnt.h), 32 bytes, no padding:
// every field here is naturally aligned at its own offset already.
type minidumpHeader struct {
	Signature          uint32
	Version            uint32
	NumberOfStreams    uint32
	StreamDirectoryRva uint32
	CheckSum           uint32
	TimeDateStamp      uint32
	Flags              uint64
}

// minidumpLocationDescriptor mirrors MINIDUMP_LOCATION_DESCRIPTOR: despite
// the name this Rva is a plain file offset, not an RVA relative to a
// loaded image -- a minidump is never "loaded", it's just read.
type minidumpLocationDescriptor struct {
	DataSize uint32
	Rva      uint32
}

// minidumpDirectory mirrors MINIDUMP_DIRECTORY, 12 bytes.
type minidumpDirectory struct {
	StreamType uint32
	Location   minidumpLocationDescriptor
}

// Minidump is a read-only handle on a Windows .dmp file, parsed directly
// from the public MINIDUMP_* structures (winnt.h) instead of calling
// dbghelp.dll's MiniDumpReadDumpStream. Same spirit as the rest of this
// package's PE/ROP parsing (see pe.go, gadgets.go, and the "reimplemented
// directly from the spec instead" note on checksec in the README): the
// format is just bytes with a documented, stable layout, and parsing it
// directly means this works without GOOS=windows or dbghelp.dll present,
// and is unit-testable against a synthetic in-memory buffer instead of
// needing a real crash dump on disk.
type Minidump struct {
	r       io.ReaderAt
	closer  io.Closer
	header  minidumpHeader
	streams []minidumpDirectory
}

// OpenMinidump opens and parses a .dmp file's header and stream directory.
func OpenMinidump(path string) (*Minidump, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	m, err := newMinidump(f)
	if err != nil {
		f.Close()
		return nil, err
	}
	m.closer = f
	return m, nil
}

// newMinidump parses from any io.ReaderAt (a file, or an in-memory
// bytes.Reader for tests/already-loaded buffers) -- OpenMinidump is just
// this plus a file open/close.
func newMinidump(r io.ReaderAt) (*Minidump, error) {
	m := &Minidump{r: r}

	if err := readStructAt(r, 0, &m.header); err != nil {
		return nil, fmt.Errorf("read MINIDUMP_HEADER: %w", err)
	}
	if m.header.Signature != minidumpSignature {
		return nil, fmt.Errorf("not a minidump file (signature 0x%X, want 0x%X)", m.header.Signature, minidumpSignature)
	}

	m.streams = make([]minidumpDirectory, m.header.NumberOfStreams)
	for i := range m.streams {
		const sizeofDirectory = 12
		off := int64(m.header.StreamDirectoryRva) + int64(i)*sizeofDirectory
		if err := readStructAt(r, off, &m.streams[i]); err != nil {
			return nil, fmt.Errorf("read MINIDUMP_DIRECTORY[%d]: %w", i, err)
		}
	}
	return m, nil
}

func (m *Minidump) Close() error {
	if m.closer != nil {
		return m.closer.Close()
	}
	return nil
}

func (m *Minidump) findStream(t MinidumpStreamType) (minidumpDirectory, bool) {
	for _, d := range m.streams {
		if d.StreamType == uint32(t) {
			return d, true
		}
	}
	return minidumpDirectory{}, false
}

// RawStream returns the raw bytes of the first stream of type t -- the
// direct analogue of MiniDumpReadDumpStream for any stream this package
// doesn't decode natively (SystemInfoStream, ThreadListStream,
// Memory64ListStream, ...). The caller is responsible for knowing that
// stream's layout.
func (m *Minidump) RawStream(t MinidumpStreamType) ([]byte, error) {
	dir, ok := m.findStream(t)
	if !ok {
		return nil, fmt.Errorf("stream type %d not present in this minidump", t)
	}
	buf := make([]byte, dir.Location.DataSize)
	if _, err := m.r.ReadAt(buf, int64(dir.Location.Rva)); err != nil {
		return nil, err
	}
	return buf, nil
}

// MinidumpModule is one entry of MINIDUMP_MODULE_LIST: a loaded module's
// name and the base address it was loaded at -- exactly what you need to
// rebase a crash address back into the binary you can actually open in a
// disassembler.
type MinidumpModule struct {
	Name          string
	BaseOfImage   uint64
	SizeOfImage   uint32
	TimeDateStamp uint32
}

// sizeofMinidumpModule is sizeof(MINIDUMP_MODULE): BaseOfImage(8) +
// SizeOfImage(4) + CheckSum(4) + TimeDateStamp(4) + ModuleNameRva(4) +
// VS_FIXEDFILEINFO(52) + CvRecord(8) + MiscRecord(8) + Reserved0(8) +
// Reserved1(8) = 108. Decoded by fixed offset below rather than a matching
// Go struct, since only a handful of its fields are useful here and
// VS_FIXEDFILEINFO's 13 DWORDs aren't worth modeling just to skip over.
const sizeofMinidumpModule = 108

// Modules walks MINIDUMP_MODULE_LIST and resolves each module's name
// string, the analogue of pwntools' Corefile module list but for a Windows
// crash dump.
func (m *Minidump) Modules() ([]MinidumpModule, error) {
	dir, ok := m.findStream(StreamModuleList)
	if !ok {
		return nil, errors.New("ModuleListStream not present in this minidump")
	}

	var count uint32
	if err := readUint32At(m.r, int64(dir.Location.Rva), &count); err != nil {
		return nil, fmt.Errorf("read MINIDUMP_MODULE_LIST.NumberOfModules: %w", err)
	}

	base := int64(dir.Location.Rva) + 4
	out := make([]MinidumpModule, 0, count)
	for i := uint32(0); i < count; i++ {
		buf := make([]byte, sizeofMinidumpModule)
		if _, err := m.r.ReadAt(buf, base+int64(i)*sizeofMinidumpModule); err != nil {
			return nil, fmt.Errorf("read MINIDUMP_MODULE[%d]: %w", i, err)
		}

		nameRva := binary.LittleEndian.Uint32(buf[20:24])
		name, err := m.readMinidumpString(nameRva)
		if err != nil {
			return nil, fmt.Errorf("read module name for MINIDUMP_MODULE[%d]: %w", i, err)
		}

		out = append(out, MinidumpModule{
			Name:          name,
			BaseOfImage:   binary.LittleEndian.Uint64(buf[0:8]),
			SizeOfImage:   binary.LittleEndian.Uint32(buf[8:12]),
			TimeDateStamp: binary.LittleEndian.Uint32(buf[16:20]),
		})
	}
	return out, nil
}

// readMinidumpString reads a MINIDUMP_STRING at the given offset: a
// ULONG32 byte length (excluding the length field and the terminator)
// followed by a UTF-16LE buffer.
func (m *Minidump) readMinidumpString(offset uint32) (string, error) {
	var length uint32
	if err := readUint32At(m.r, int64(offset), &length); err != nil {
		return "", err
	}
	buf := make([]byte, length)
	if _, err := m.r.ReadAt(buf, int64(offset)+4); err != nil {
		return "", err
	}
	units := make([]uint16, length/2)
	for i := range units {
		units[i] = binary.LittleEndian.Uint16(buf[i*2:])
	}
	return string(utf16.Decode(units)), nil
}

// exceptionMaxParameters is EXCEPTION_MAXIMUM_PARAMETERS (winnt.h): the
// fixed size of MINIDUMP_EXCEPTION.ExceptionInformation.
const exceptionMaxParameters = 15

// MinidumpException is MINIDUMP_EXCEPTION_STREAM flattened to the fields a
// crash-triage script actually wants: which thread, what kind of fault
// (ExceptionCode -- e.g. 0xC0000005 for an access violation, the same
// value Tube.Interactive already reports for a locally observed crash),
// and where.
type MinidumpException struct {
	ThreadID         uint32
	ExceptionCode    uint32
	ExceptionFlags   uint32
	ExceptionAddress uint64
	// Parameters holds the first NumberParameters entries of
	// ExceptionInformation -- e.g. for an access violation, Parameters[0]
	// is the access type (read/write/execute) and Parameters[1] is the
	// faulting address.
	Parameters []uint64
}

// sizeofMinidumpExceptionStream is sizeof(MINIDUMP_EXCEPTION_STREAM):
// ThreadId(4) + alignment(4) + MINIDUMP_EXCEPTION(152) +
// ThreadContext location descriptor(8) = 168.
const sizeofMinidumpExceptionStream = 168

// Exception decodes MINIDUMP_EXCEPTION_STREAM, if present (a minidump
// taken from a still-running, non-crashed process has no exception
// stream). The register context blob referenced by
// MINIDUMP_EXCEPTION_STREAM.ThreadContext is not decoded here -- CONTEXT's
// layout differs by architecture and has internal padding/XSAVE-area
// subtleties not worth getting wrong; use RawStream(StreamException) and
// slice past sizeofMinidumpExceptionStream's ThreadContext location if you
// need the raw register bytes for a specific architecture.
func (m *Minidump) Exception() (*MinidumpException, error) {
	dir, ok := m.findStream(StreamException)
	if !ok {
		return nil, errors.New("ExceptionStream not present in this minidump (the process may not have crashed)")
	}

	buf := make([]byte, sizeofMinidumpExceptionStream)
	if _, err := m.r.ReadAt(buf, int64(dir.Location.Rva)); err != nil {
		return nil, fmt.Errorf("read MINIDUMP_EXCEPTION_STREAM: %w", err)
	}

	threadID := binary.LittleEndian.Uint32(buf[0:4])
	// MINIDUMP_EXCEPTION starts right after ThreadId + a 4-byte alignment pad.
	exc := buf[8:]
	numParams := binary.LittleEndian.Uint32(exc[24:28])
	if numParams > exceptionMaxParameters {
		numParams = exceptionMaxParameters
	}
	params := make([]uint64, numParams)
	for i := range params {
		params[i] = binary.LittleEndian.Uint64(exc[32+i*8:])
	}

	return &MinidumpException{
		ThreadID:         threadID,
		ExceptionCode:    binary.LittleEndian.Uint32(exc[0:4]),
		ExceptionFlags:   binary.LittleEndian.Uint32(exc[4:8]),
		ExceptionAddress: binary.LittleEndian.Uint64(exc[16:24]),
		Parameters:       params,
	}, nil
}

// readStructAt fills v (a pointer to a fixed-size struct of fixed-width
// fields) by reading binary.Size(v) bytes at offset -- the minidump.go
// analogue of PEFile.readStructAt in pe.go, kept separate since Minidump
// isn't a PEFile and has no reason to share its receiver.
func readStructAt(r io.ReaderAt, offset int64, v any) error {
	size := binary.Size(v)
	if size < 0 {
		return errors.New("readStructAt: unsupported type")
	}
	buf := make([]byte, size)
	if _, err := r.ReadAt(buf, offset); err != nil {
		return err
	}
	return binary.Read(bytes.NewReader(buf), binary.LittleEndian, v)
}

func readUint32At(r io.ReaderAt, offset int64, out *uint32) error {
	var buf [4]byte
	if _, err := r.ReadAt(buf[:], offset); err != nil {
		return err
	}
	*out = binary.LittleEndian.Uint32(buf[:])
	return nil
}