package infra

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"os"
)

const (
	rtIcon      = 3
	rtGroupIcon = 14

	scnCntInitializedData = 0x00000040
	scnMemRead            = 0x40000000

	dataDirResource = 2
	resourceTableSlot = 2
)

type icoEntry struct {
	width    uint8
	height   uint8
	colors   uint8
	reserved uint8
	planes   uint16
	bitCount uint16
	size     uint32
	offset   uint32
	data     []byte
}

func parseICO(buf []byte) ([]icoEntry, error) {
	if len(buf) < 6 {
		return nil, fmt.Errorf("ico too short")
	}
	if binary.LittleEndian.Uint16(buf[0:2]) != 0 {
		return nil, fmt.Errorf("ico reserved field must be 0")
	}
	if binary.LittleEndian.Uint16(buf[2:4]) != 1 {
		return nil, fmt.Errorf("not an icon (type != 1)")
	}
	count := int(binary.LittleEndian.Uint16(buf[4:6]))
	if count == 0 {
		return nil, fmt.Errorf("ico has 0 images")
	}
	if len(buf) < 6+16*count {
		return nil, fmt.Errorf("ico header truncated")
	}
	entries := make([]icoEntry, count)
	for i := 0; i < count; i++ {
		off := 6 + 16*i
		e := icoEntry{
			width:    buf[off],
			height:   buf[off+1],
			colors:   buf[off+2],
			reserved: buf[off+3],
			planes:   binary.LittleEndian.Uint16(buf[off+4 : off+6]),
			bitCount: binary.LittleEndian.Uint16(buf[off+6 : off+8]),
			size:     binary.LittleEndian.Uint32(buf[off+8 : off+12]),
			offset:   binary.LittleEndian.Uint32(buf[off+12 : off+16]),
		}
		if int(e.offset)+int(e.size) > len(buf) {
			return nil, fmt.Errorf("ico image %d out of bounds", i)
		}
		e.data = buf[e.offset : e.offset+e.size]
		entries[i] = e
	}
	return entries, nil
}

type peLayout struct {
	data           []byte
	peOff          int
	is64           bool
	numSections    int
	sizeOptHdr     int
	sectionAlign   uint32
	fileAlign      uint32
	sectionsStart  int
	sizeOfImage    uint32
	dataDirCount   uint32
	dataDirOff     int
	checksumOff    int
	sizeOfImageOff int
}

func parsePE(buf []byte) (*peLayout, error) {
	if len(buf) < 64 || buf[0] != 'M' || buf[1] != 'Z' {
		return nil, fmt.Errorf("not a PE file (no MZ signature)")
	}
	peOff := int(binary.LittleEndian.Uint32(buf[0x3C:0x40]))
	if peOff+24 > len(buf) {
		return nil, fmt.Errorf("PE header out of bounds")
	}
	if string(buf[peOff:peOff+4]) != "PE\x00\x00" {
		return nil, fmt.Errorf("not a PE file (no PE signature at e_lfanew)")
	}
	coff := peOff + 4
	numSections := int(binary.LittleEndian.Uint16(buf[coff+2 : coff+4]))
	sizeOptHdr := int(binary.LittleEndian.Uint16(buf[coff+16 : coff+18]))
	optOff := coff + 20
	if optOff+sizeOptHdr > len(buf) {
		return nil, fmt.Errorf("optional header truncated")
	}
	magic := binary.LittleEndian.Uint16(buf[optOff : optOff+2])
	var is64 bool
	switch magic {
	case 0x10B:
		is64 = false
	case 0x20B:
		is64 = true
	default:
		return nil, fmt.Errorf("unknown optional header magic 0x%X", magic)
	}

	var sectionAlignOff, fileAlignOff, sizeOfImageOff, checksumOff, numRvaOff, dataDirOff int
	if is64 {
		sectionAlignOff = optOff + 32
		fileAlignOff = optOff + 36
		sizeOfImageOff = optOff + 56
		checksumOff = optOff + 64
		numRvaOff = optOff + 108
		dataDirOff = optOff + 112
	} else {
		sectionAlignOff = optOff + 32
		fileAlignOff = optOff + 36
		sizeOfImageOff = optOff + 56
		checksumOff = optOff + 64
		numRvaOff = optOff + 92
		dataDirOff = optOff + 96
	}
	dataDirCount := binary.LittleEndian.Uint32(buf[numRvaOff : numRvaOff+4])

	return &peLayout{
		data:           buf,
		peOff:          peOff,
		is64:           is64,
		numSections:    numSections,
		sizeOptHdr:     sizeOptHdr,
		sectionAlign:   binary.LittleEndian.Uint32(buf[sectionAlignOff : sectionAlignOff+4]),
		fileAlign:      binary.LittleEndian.Uint32(buf[fileAlignOff : fileAlignOff+4]),
		sectionsStart:  optOff + sizeOptHdr,
		sizeOfImage:    binary.LittleEndian.Uint32(buf[sizeOfImageOff : sizeOfImageOff+4]),
		dataDirCount:   dataDirCount,
		dataDirOff:     dataDirOff,
		checksumOff:    checksumOff,
		sizeOfImageOff: sizeOfImageOff,
	}, nil
}

type sectionHdr struct {
	name             string
	virtualSize      uint32
	virtualAddress   uint32
	sizeOfRawData    uint32
	pointerToRawData uint32
	characteristics  uint32
	headerOff        int
}

func (p *peLayout) sections() []sectionHdr {
	out := make([]sectionHdr, p.numSections)
	for i := 0; i < p.numSections; i++ {
		off := p.sectionsStart + 40*i
		nameBytes := p.data[off : off+8]
		end := bytes.IndexByte(nameBytes, 0)
		if end < 0 {
			end = 8
		}
		out[i] = sectionHdr{
			name:             string(nameBytes[:end]),
			virtualSize:      binary.LittleEndian.Uint32(p.data[off+8 : off+12]),
			virtualAddress:   binary.LittleEndian.Uint32(p.data[off+12 : off+16]),
			sizeOfRawData:    binary.LittleEndian.Uint32(p.data[off+16 : off+20]),
			pointerToRawData: binary.LittleEndian.Uint32(p.data[off+20 : off+24]),
			characteristics:  binary.LittleEndian.Uint32(p.data[off+36 : off+40]),
			headerOff:        off,
		}
	}
	return out
}

func (p *peLayout) hasRsrc() bool {
	if p.dataDirCount > resourceTableSlot {
		off := p.dataDirOff + resourceTableSlot*8
		va := binary.LittleEndian.Uint32(p.data[off : off+4])
		size := binary.LittleEndian.Uint32(p.data[off+4 : off+8])
		if va != 0 && size != 0 {
			return true
		}
	}
	for _, s := range p.sections() {
		if s.name == ".rsrc" {
			return true
		}
	}
	return false
}

func alignUp(v, a uint32) uint32 {
	if a == 0 {
		return v
	}
	return (v + a - 1) &^ (a - 1)
}

func buildResourceSection(entries []icoEntry, baseRVA uint32) ([]byte, error) {
	n := uint32(len(entries))

	rootDirSize := uint32(16 + 2*8)
	rtIconDirSize := uint32(16 + n*8)
	perIconNameDirSize := uint32(16 + 8)
	rtGroupDirSize := uint32(16 + 8)
	groupNameDirSize := uint32(16 + 8)

	leafEntrySize := uint32(16)
	totalLeafEntries := n + 1
	leavesSize := leafEntrySize * totalLeafEntries

	dirsSize := rootDirSize + rtIconDirSize + n*perIconNameDirSize + rtGroupDirSize + groupNameDirSize

	dirsAndLeaves := dirsSize + leavesSize

	groupDataSize := uint32(6 + 14*n)
	dataStartOff := dirsAndLeaves
	groupDataOff := dataStartOff
	groupDataPadded := alignUp(groupDataSize, 4)

	iconDataOffsets := make([]uint32, n)
	cursor := groupDataOff + groupDataPadded
	for i, e := range entries {
		iconDataOffsets[i] = cursor
		cursor += alignUp(uint32(len(e.data)), 4)
	}
	totalSize := cursor

	out := make([]byte, totalSize)

	leafOff := dirsSize
	groupLeafOff := leafOff
	iconLeavesStart := leafOff + leafEntrySize

	rootOff := uint32(0)
	rtGroupSubOff := rootDirSize
	rtIconSubOff := rtGroupSubOff + rtGroupDirSize
	groupNameSubOff := rtIconSubOff + rtIconDirSize
	perIconNamesStart := groupNameSubOff + groupNameDirSize

	writeDirHeader := func(off uint32, idEntries uint16) {
		binary.LittleEndian.PutUint32(out[off:off+4], 0)
		binary.LittleEndian.PutUint32(out[off+4:off+8], 0)
		binary.LittleEndian.PutUint16(out[off+8:off+10], 0)
		binary.LittleEndian.PutUint16(out[off+10:off+12], 0)
		binary.LittleEndian.PutUint16(out[off+12:off+14], 0)
		binary.LittleEndian.PutUint16(out[off+14:off+16], idEntries)
	}
	writeIDEntry := func(off uint32, id uint32, target uint32, isDir bool) {
		binary.LittleEndian.PutUint32(out[off:off+4], id)
		val := target
		if isDir {
			val |= 0x80000000
		}
		binary.LittleEndian.PutUint32(out[off+4:off+8], val)
	}

	writeDirHeader(rootOff, 2)
	writeIDEntry(rootOff+16, rtIcon, rtIconSubOff, true)
	writeIDEntry(rootOff+24, rtGroupIcon, rtGroupSubOff, true)

	writeDirHeader(rtGroupSubOff, 1)
	writeIDEntry(rtGroupSubOff+16, 1, groupNameSubOff, true)

	writeDirHeader(groupNameSubOff, 1)
	writeIDEntry(groupNameSubOff+16, 0, groupLeafOff, false)

	writeDirHeader(rtIconSubOff, uint16(n))
	for i := uint32(0); i < n; i++ {
		writeIDEntry(rtIconSubOff+16+i*8, i+1, perIconNamesStart+i*perIconNameDirSize, true)
	}
	for i := uint32(0); i < n; i++ {
		nameDirOff := perIconNamesStart + i*perIconNameDirSize
		writeDirHeader(nameDirOff, 1)
		leafForIcon := iconLeavesStart + i*leafEntrySize
		writeIDEntry(nameDirOff+16, 0, leafForIcon, false)
	}

	writeLeaf := func(off uint32, dataOff uint32, size uint32) {
		binary.LittleEndian.PutUint32(out[off:off+4], baseRVA+dataOff)
		binary.LittleEndian.PutUint32(out[off+4:off+8], size)
		binary.LittleEndian.PutUint32(out[off+8:off+12], 0)
		binary.LittleEndian.PutUint32(out[off+12:off+16], 0)
	}

	writeLeaf(groupLeafOff, groupDataOff, groupDataSize)
	for i := uint32(0); i < n; i++ {
		writeLeaf(iconLeavesStart+i*leafEntrySize, iconDataOffsets[i], uint32(len(entries[i].data)))
	}

	binary.LittleEndian.PutUint16(out[groupDataOff:groupDataOff+2], 0)
	binary.LittleEndian.PutUint16(out[groupDataOff+2:groupDataOff+4], 1)
	binary.LittleEndian.PutUint16(out[groupDataOff+4:groupDataOff+6], uint16(n))
	for i, e := range entries {
		eo := groupDataOff + 6 + uint32(i)*14
		out[eo] = e.width
		out[eo+1] = e.height
		out[eo+2] = e.colors
		out[eo+3] = e.reserved
		binary.LittleEndian.PutUint16(out[eo+4:eo+6], e.planes)
		binary.LittleEndian.PutUint16(out[eo+6:eo+8], e.bitCount)
		binary.LittleEndian.PutUint32(out[eo+8:eo+12], e.size)
		binary.LittleEndian.PutUint16(out[eo+12:eo+14], uint16(i+1))
	}

	for i, e := range entries {
		copy(out[iconDataOffsets[i]:], e.data)
	}

	return out, nil
}

// SetExeIcon embebe el icono del archivo .ico en el .exe sobreescribiendo
// el archivo. Funciona para PE32 y PE32+ que aun no tienen seccion .rsrc
// (caso comun de binarios Go compilados sin icono). Si el .exe ya tiene
// recursos retorna error.
func SetExeIcon(exePath, icoPath string) error {
	icoBuf, err := os.ReadFile(icoPath)
	if err != nil {
		return fmt.Errorf("read ico: %w", err)
	}
	entries, err := parseICO(icoBuf)
	if err != nil {
		return fmt.Errorf("parse ico: %w", err)
	}

	exeBuf, err := os.ReadFile(exePath)
	if err != nil {
		return fmt.Errorf("read exe: %w", err)
	}
	pe, err := parsePE(exeBuf)
	if err != nil {
		return fmt.Errorf("parse pe: %w", err)
	}
	if pe.hasRsrc() {
		return fmt.Errorf("exe already has .rsrc resources; not supported")
	}
	if pe.dataDirCount <= resourceTableSlot {
		return fmt.Errorf("optional header DataDirectory has only %d entries (need >= %d)", pe.dataDirCount, resourceTableSlot+1)
	}

	sections := pe.sections()
	if len(sections) == 0 {
		return fmt.Errorf("exe has no sections")
	}
	last := sections[len(sections)-1]

	newSecHdrOff := pe.sectionsStart + 40*pe.numSections
	if newSecHdrOff+40 > int(last.pointerToRawData) {
		return fmt.Errorf("not enough space in PE headers for new section header")
	}

	newRVA := alignUp(last.virtualAddress+last.virtualSize, pe.sectionAlign)
	newRawOff := alignUp(last.pointerToRawData+last.sizeOfRawData, pe.fileAlign)

	rsrc, err := buildResourceSection(entries, newRVA)
	if err != nil {
		return fmt.Errorf("build resource section: %w", err)
	}
	rawSize := alignUp(uint32(len(rsrc)), pe.fileAlign)
	virtSize := uint32(len(rsrc))

	out := make([]byte, 0, int(newRawOff)+int(rawSize))
	out = append(out, exeBuf[:newRawOff]...)
	if int(newRawOff) > len(exeBuf) {
		out = append(out, make([]byte, int(newRawOff)-len(exeBuf))...)
	}
	out = append(out, rsrc...)
	if rawSize > uint32(len(rsrc)) {
		out = append(out, make([]byte, rawSize-uint32(len(rsrc)))...)
	}

	hdr := make([]byte, 40)
	copy(hdr[0:8], ".rsrc\x00\x00\x00")
	binary.LittleEndian.PutUint32(hdr[8:12], virtSize)
	binary.LittleEndian.PutUint32(hdr[12:16], newRVA)
	binary.LittleEndian.PutUint32(hdr[16:20], rawSize)
	binary.LittleEndian.PutUint32(hdr[20:24], newRawOff)
	binary.LittleEndian.PutUint32(hdr[36:40], scnCntInitializedData|scnMemRead)
	copy(out[newSecHdrOff:newSecHdrOff+40], hdr)

	binary.LittleEndian.PutUint16(out[pe.peOff+4+2:pe.peOff+4+4], uint16(pe.numSections+1))

	rsrcEntryOff := pe.dataDirOff + resourceTableSlot*8
	binary.LittleEndian.PutUint32(out[rsrcEntryOff:rsrcEntryOff+4], newRVA)
	binary.LittleEndian.PutUint32(out[rsrcEntryOff+4:rsrcEntryOff+8], virtSize)

	newSizeOfImage := alignUp(newRVA+virtSize, pe.sectionAlign)
	binary.LittleEndian.PutUint32(out[pe.sizeOfImageOff:pe.sizeOfImageOff+4], newSizeOfImage)

	binary.LittleEndian.PutUint32(out[pe.checksumOff:pe.checksumOff+4], 0)

	if err := os.WriteFile(exePath, out, 0o755); err != nil {
		return fmt.Errorf("write exe: %w", err)
	}
	return nil
}