package pe

import (
	"bytes"
	"encoding/binary"
	"io"

	"git.akyoto.dev/cli/q/src/config"
	"git.akyoto.dev/cli/q/src/exe"
)

// EXE is the portable executable format used on Windows.
type EXE struct {
	DOSHeader
	NTHeader
	OptionalHeader64
	Sections []SectionHeader
}

// Write writes the EXE file to the given writer.
func Write(writer io.Writer, code []byte, data []byte, dlls []DLL) {
	NumSections := 2
	HeaderEnd := DOSHeaderSize + NTHeaderSize + OptionalHeader64Size + SectionHeaderSize*NumSections
	codeStart, codePadding := exe.Align(HeaderEnd, config.Align)
	dataStart, dataPadding := exe.Align(codeStart+len(code), config.Align)

	dllImports := []DLLImport{}

	for _, dll := range dlls {
		dllAddresses := []uint64{}
		dllNamePos := len(data)
		data = append(data, dll.Name...)
		data = append(data, 0x00)

		for _, f := range dll.Functions {
			pos := len(data)
			data = append(data, 0x00, 0x00)
			data = append(data, f...)
			data = append(data, 0x00)

			if len(data)&1 != 0 {
				data = append(data, 0x00) // align the next entry on an even boundary
			}

			dllAddresses = append(dllAddresses, uint64(dataStart+pos))
		}

		dllAddresses = append(dllAddresses, 0)

		// Add the address table to the data section
		functionAddressesStart := dataStart + len(data)
		data, _ = binary.Append(data, binary.LittleEndian, &dllAddresses)

		dllImports = append(dllImports, DLLImport{
			RvaFunctionNameList:    uint32(functionAddressesStart),
			TimeDateStamp:          0,
			ForwarderChain:         0,
			RvaModuleName:          uint32(dataStart + dllNamePos),
			RvaFunctionAddressList: uint32(functionAddressesStart),
		})
	}

	dllImports = append(dllImports, DLLImport{}) // a zeroed structure marks the end of the list

	// Add imports to the data section
	importsStart := dataStart + len(data)
	importsSize := DLLImportSize * len(dllImports)
	data, _ = binary.Append(data, binary.LittleEndian, &dllImports)

	imageSize := dataStart + len(data)
	imageSize, _ = exe.Align(imageSize, config.Align)

	pe := &EXE{
		DOSHeader: DOSHeader{
			Magic:          [4]byte{'M', 'Z', 0, 0},
			NTHeaderOffset: DOSHeaderSize,
		},
		NTHeader: NTHeader{
			Signature:            [4]byte{'P', 'E', 0, 0},
			Machine:              IMAGE_FILE_MACHINE_AMD64,
			NumberOfSections:     uint16(NumSections),
			TimeDateStamp:        0,
			PointerToSymbolTable: 0,
			NumberOfSymbols:      0,
			SizeOfOptionalHeader: OptionalHeader64Size,
			Characteristics:      IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_LARGE_ADDRESS_AWARE,
		},
		OptionalHeader64: OptionalHeader64{
			Magic:                       0x020B, // PE32+ / 64-bit executable
			MajorLinkerVersion:          0x0E,
			MinorLinkerVersion:          0x16,
			SizeOfCode:                  uint32(len(code)),
			SizeOfInitializedData:       0,
			SizeOfUninitializedData:     0,
			AddressOfEntryPoint:         config.CodeOffset,
			BaseOfCode:                  config.CodeOffset,
			ImageBase:                   config.BaseAddress,
			SectionAlignment:            config.Align, // power of 2, must be greater than or equal to FileAlignment
			FileAlignment:               config.Align, // power of 2
			MajorOperatingSystemVersion: 0x06,
			MinorOperatingSystemVersion: 0,
			MajorImageVersion:           0,
			MinorImageVersion:           0,
			MajorSubsystemVersion:       0x06,
			MinorSubsystemVersion:       0,
			Win32VersionValue:           0,
			SizeOfImage:                 uint32(imageSize),
			SizeOfHeaders:               config.CodeOffset, // section bodies begin here
			CheckSum:                    0,
			Subsystem:                   IMAGE_SUBSYSTEM_WINDOWS_CUI,
			DllCharacteristics:          IMAGE_DLLCHARACTERISTICS_HIGH_ENTROPY_VA | IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE | IMAGE_DLLCHARACTERISTICS_NX_COMPAT | IMAGE_DLLCHARACTERISTICS_TERMINAL_SERVER_AWARE,
			SizeOfStackReserve:          0x100000,
			SizeOfStackCommit:           0x1000,
			SizeOfHeapReserve:           0x100000,
			SizeOfHeapCommit:            0x1000,
			LoaderFlags:                 0,
			NumberOfRvaAndSizes:         16,
			DataDirectory: [16]DataDirectory{
				{VirtualAddress: 0, Size: 0},
				{VirtualAddress: uint32(importsStart), Size: uint32(importsSize)}, // RVA of the imported function table
				{VirtualAddress: 0, Size: 0},
				{VirtualAddress: 0, Size: 0},
				{VirtualAddress: 0, Size: 0},
				{VirtualAddress: 0, Size: 0},
				{VirtualAddress: 0, Size: 0},
				{VirtualAddress: 0, Size: 0},
				{VirtualAddress: 0, Size: 0},
				{VirtualAddress: 0, Size: 0},
				{VirtualAddress: 0, Size: 0},
				{VirtualAddress: 0, Size: 0},
				{VirtualAddress: 0, Size: 0},
				{VirtualAddress: 0, Size: 0},
				{VirtualAddress: 0, Size: 0},
				{VirtualAddress: 0, Size: 0},
			},
		},
		Sections: []SectionHeader{
			{
				Name:            [8]byte{'.', 't', 'e', 'x', 't'},
				VirtualSize:     uint32(len(code)),
				VirtualAddress:  config.CodeOffset,
				RawSize:         uint32(len(code)), // must be a multiple of FileAlignment
				RawAddress:      config.CodeOffset, // must be a multiple of FileAlignment
				Characteristics: IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ,
			},
			{
				Name:            [8]byte{'.', 'r', 'd', 'a', 't', 'a'},
				VirtualSize:     uint32(len(data)),
				VirtualAddress:  uint32(dataStart),
				RawSize:         uint32(len(data)),
				RawAddress:      uint32(dataStart),
				Characteristics: IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ,
			},
		},
	}

	binary.Write(writer, binary.LittleEndian, &pe.DOSHeader)
	binary.Write(writer, binary.LittleEndian, &pe.NTHeader)
	binary.Write(writer, binary.LittleEndian, &pe.OptionalHeader64)
	binary.Write(writer, binary.LittleEndian, &pe.Sections)

	writer.Write(bytes.Repeat([]byte{0x00}, int(codePadding)))
	writer.Write(code)
	writer.Write(bytes.Repeat([]byte{0x00}, int(dataPadding)))
	writer.Write(data)
}