package pe

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

	"git.akyoto.dev/cli/q/src/config"
	"git.akyoto.dev/cli/q/src/dll"
	"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.List) {
	if len(data) == 0 {
		data = []byte{0}
	}

	NumSections := 3
	HeaderEnd := DOSHeaderSize + NTHeaderSize + OptionalHeader64Size + SectionHeaderSize*NumSections
	codeStart, codePadding := exe.Align(HeaderEnd, config.Align)
	dataStart, dataPadding := exe.Align(codeStart+len(code), config.Align)
	importsStart, importsPadding := exe.Align(dataStart+len(data), config.Align)

	subSystem := IMAGE_SUBSYSTEM_WINDOWS_CUI

	if dlls.Contains("user32") {
		subSystem = IMAGE_SUBSYSTEM_WINDOWS_GUI
	}

	imports := make([]uint64, 0)
	dllData := make([]byte, 0)
	dllImports := []DLLImport{}

	for _, library := range dlls {
		functionsStart := len(imports) * 8
		dllNamePos := len(dllData)
		dllData = append(dllData, library.Name...)
		dllData = append(dllData, ".dll"...)
		dllData = append(dllData, 0x00)

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

		for _, fn := range library.Functions {
			if len(dllData)&1 != 0 {
				dllData = append(dllData, 0x00) // align the next entry on an even boundary
			}

			offset := len(dllData)
			dllData = append(dllData, 0x00, 0x00)
			dllData = append(dllData, fn...)
			dllData = append(dllData, 0x00)

			imports = append(imports, uint64(offset))
		}

		imports = append(imports, 0)
	}

	dllDataStart := importsStart + len(imports)*8

	for i := range imports {
		if imports[i] == 0 {
			continue
		}

		imports[i] += uint64(dllDataStart)
	}

	for i := range dllImports {
		dllImports[i].RvaModuleName += uint32(dllDataStart)
	}

	dllImports = append(dllImports, DLLImport{}) // a zeroed structure marks the end of the list
	importDirectoryStart := dllDataStart + len(dllData)
	importDirectorySize := DLLImportSize * len(dllImports)

	importSectionSize := len(imports)*8 + len(dllData) + importDirectorySize
	imageSize := importsStart + importSectionSize
	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:                   uint16(subSystem),
			DllCharacteristics:          IMAGE_DLLCHARACTERISTICS_HIGH_ENTROPY_VA | IMAGE_DLLCHARACTERISTICS_NX_COMPAT | IMAGE_DLLCHARACTERISTICS_TERMINAL_SERVER_AWARE, // IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE
			SizeOfStackReserve:          0x100000,
			SizeOfStackCommit:           0x1000,
			SizeOfHeapReserve:           0x100000,
			SizeOfHeapCommit:            0x1000,
			LoaderFlags:                 0,
			NumberOfRvaAndSizes:         16,
			DataDirectory: [16]DataDirectory{
				{VirtualAddress: 0, Size: 0},
				{VirtualAddress: uint32(importDirectoryStart), Size: uint32(importDirectorySize)}, // 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: uint32(importsStart), Size: uint32(len(imports) * 8)}, // RVA of the import address table
				{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,
			},
			{
				Name:            [8]byte{'.', 'i', 'd', 'a', 't', 'a'},
				VirtualSize:     uint32(importSectionSize),
				VirtualAddress:  uint32(importsStart),
				RawSize:         uint32(importSectionSize),
				RawAddress:      uint32(importsStart),
				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)
	writer.Write(bytes.Repeat([]byte{0x00}, int(importsPadding)))
	binary.Write(writer, binary.LittleEndian, &imports)
	binary.Write(writer, binary.LittleEndian, &dllData)
	binary.Write(writer, binary.LittleEndian, &dllImports)
}