package image

import (
	"fmt"
	"io"
	"os"
)

const pxeInitScript = `#!/bin/sh
mount -t proc proc /proc 2>/dev/null
mount -t sysfs sys /sys 2>/dev/null
mount -t devtmpfs dev /dev 2>/dev/null
mkdir -p /tmp

ISO_URL=""
for param in $(cat /proc/cmdline); do
    case $param in
        prov.iso=*) ISO_URL="${param#prov.iso=}" ;;
    esac
done

if [ -n "$ISO_URL" ]; then
    echo "=== PXE: Setting up network ==="
    for mod in e1000 e1000e igb ixgbe i40e ice tg3 bnxt_en r8169 virtio_net; do
        modprobe $mod 2>/dev/null
    done
    sleep 2

    for iface in $(ls /sys/class/net/ 2>/dev/null | grep -v lo); do
        ip link set "$iface" up 2>/dev/null
        if udhcpc -i "$iface" -n -q -t 5 2>/dev/null; then
            echo "PXE: Network up on $iface"
            break
        fi
    done

    echo "=== PXE: Downloading ISO ==="
    wget -O /tmp/proxmox.iso "$ISO_URL" 2>&1

    if [ -s /tmp/proxmox.iso ]; then
        losetup /dev/loop0 /tmp/proxmox.iso
        echo "PXE: ISO on /dev/loop0"
    else
        echo "PXE: ERROR - download failed"
    fi
fi

umount /dev 2>/dev/null
umount /sys 2>/dev/null
umount /proc 2>/dev/null
exec /init "$@"
`

func CreatePXEOverlay(outputPath string) error {
	out, err := os.Create(outputPath)
	if err != nil {
		return fmt.Errorf("create overlay: %w", err)
	}
	defer out.Close()

	if err := writeCPIOEntry(out, "pxe-init", []byte(pxeInitScript), 0o100755, 1); err != nil {
		return err
	}
	return writeCPIOTrailer(out)
}

func writeCPIOEntry(w io.Writer, name string, data []byte, mode uint32, ino uint32) error {
	if err := writeCPIOHeader(w, name, int64(len(data)), mode, ino); err != nil {
		return err
	}
	if _, err := w.Write(data); err != nil {
		return err
	}
	if pad := (4 - len(data)%4) % 4; pad > 0 {
		if _, err := w.Write(make([]byte, pad)); err != nil {
			return err
		}
	}
	return nil
}

func writeCPIOHeader(w io.Writer, name string, fileSize int64, mode uint32, ino uint32) error {
	nameLen := len(name) + 1
	hdr := fmt.Sprintf("070701%08X%08X%08X%08X%08X%08X%08X%08X%08X%08X%08X%08X%08X",
		ino,      // ino
		mode,     // mode
		0,        // uid
		0,        // gid
		1,        // nlink
		0,        // mtime
		fileSize, // filesize
		0,        // devmajor
		0,        // devminor
		0,        // rdevmajor
		0,        // rdevminor
		nameLen,  // namesize
		0,        // check
	)
	if _, err := io.WriteString(w, hdr); err != nil {
		return err
	}
	if _, err := io.WriteString(w, name); err != nil {
		return err
	}
	if _, err := w.Write([]byte{0}); err != nil {
		return err
	}
	totalHeader := 110 + nameLen
	if pad := (4 - totalHeader%4) % 4; pad > 0 {
		if _, err := w.Write(make([]byte, pad)); err != nil {
			return err
		}
	}
	return nil
}

func writeCPIOTrailer(w io.Writer) error {
	return writeCPIOHeader(w, "TRAILER!!!", 0, 0, 0)
}