package image

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

const pxeInitScript = `#!/bin/sh
mkdir -p /dev
mount -t devtmpfs devtmpfs /dev 2>/dev/null
losetup /dev/loop0 /proxmox.iso 2>/dev/null
umount /dev 2>/dev/null
exec /init "$@"
`

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

	orig, err := os.Open(origInitrdPath)
	if err != nil {
		return fmt.Errorf("open initrd: %w", err)
	}
	if _, err := io.Copy(out, orig); err != nil {
		orig.Close()
		return fmt.Errorf("copy initrd: %w", err)
	}
	orig.Close()

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

	isoStat, err := os.Stat(isoPath)
	if err != nil {
		return fmt.Errorf("stat iso: %w", err)
	}
	isoFile, err := os.Open(isoPath)
	if err != nil {
		return fmt.Errorf("open iso: %w", err)
	}
	defer isoFile.Close()
	if err := writeCPIOEntryFromReader(out, "proxmox.iso", isoFile, isoStat.Size(), 0o100644, 2); 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 writeCPIOEntryFromReader(w io.Writer, name string, r io.Reader, size int64, mode uint32, ino uint32) error {
	if err := writeCPIOHeader(w, name, size, mode, ino); err != nil {
		return err
	}
	if _, err := io.Copy(w, r); err != nil {
		return err
	}
	if pad := (4 - int(size%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)
}