package main

import (
	"archive/tar"
	"bytes"
	"debug/pe"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"time"
)

func rmTempFile(f *os.File) {
	f.Close()
	if err := os.Remove(f.Name()); err != nil {
		log.Print("failed to remove ", f.Name(), ": ", err)
	}
}

func buildUki(out io.Writer, ctx *renderContext, uki bool, cmdline string) error {
	if !uki {
		return fmt.Errorf("buildUki only builds UKI")
	}

	const efiStubPath = "/usr/lib/systemd/boot/efi/linuxx64.efi.stub"

	stub, err := pe.Open(efiStubPath)
	if err != nil {
		return fmt.Errorf("efi stub open: %w", err)
	}

	hdr := stub.OptionalHeader.(*pe.OptionalHeader64)
	align := uint64(hdr.SectionAlignment)

	var offset uint64
	for _, s := range stub.Sections {
		end := uint64(s.VirtualAddress) + uint64(s.VirtualSize)
		if end > offset {
			offset = end
		}
	}

	offset += hdr.ImageBase

	stub.Close()

	// kernel
	kernelPath, err := distFetch("kernels", ctx.Host.Kernel)
	if err != nil {
		return fmt.Errorf("fetch kernel: %w", err)
	}

	// temp dir
	tempDir, err := os.MkdirTemp(os.TempDir(), "uki-")
	if err != nil {
		return fmt.Errorf("mkdir temp: %w", err)
	}
	defer os.RemoveAll(tempDir)

	// osrel
	osrelPath := filepath.Join(tempDir, "osrel")
	if err := os.WriteFile(osrelPath, fmt.Appendf(nil, "ID=direktil\nPRETTY_NAME='Direktil %s %s+0'\n", ctx.Host.Name,
		time.Now().UTC().Format(time.DateTime)), 0o600); err != nil {
		return fmt.Errorf("create osrel: %w", err)
	}

	// cmdline
	cmdlinePath := filepath.Join(tempDir, "cmdline")
	if err := os.WriteFile(cmdlinePath, append([]byte(cmdline), 0), 0o600); err != nil {
		return fmt.Errorf("create cmdline: %w", err)
	}

	// initrd
	initrdPath := filepath.Join(tempDir, "initrd")
	if err := func() (err error) {
		initrd, err := os.Create(initrdPath)
		if err != nil {
			return
		}
		defer initrd.Close()
		return buildInitrd(initrd, ctx)
	}(); err != nil {
		return fmt.Errorf("create initrd: %w", err)
	}

	// assemble

	args := make([]string, 0)
	for _, i := range []struct {
		section string
		path    string
	}{
		{"osrel", osrelPath},
		{"cmdline", cmdlinePath},
		{"initrd", initrdPath},
		{"linux", kernelPath},
	} {
		offset += align - offset%align
		args = append(args,
			"--add-section", "."+i.section+"="+i.path, "--change-section-vma", fmt.Sprintf(".%s=0x%x", i.section, offset))

		stat, err := os.Stat(i.path)
		if err != nil {
			return fmt.Errorf("stat %s: %w", i.section, err)
		}

		offset += uint64(stat.Size())
	}

	ukiPath := filepath.Join(tempDir, "uki")
	args = append(args, efiStubPath, ukiPath)

	if err := run("objcopy", args...); err != nil {
		return fmt.Errorf("objcopy: %w", err)
	}

	// read
	ukiBytes, err := os.ReadFile(ukiPath)
	if err != nil {
		return fmt.Errorf("read uki: %w", err)
	}

	io.Copy(out, bytes.NewBuffer(ukiBytes))

	// done
	return nil
}

func buildBootTar(out io.Writer, ctx *renderContext, uki bool, cmdline string) (err error) {
	if uki {
		return buildUkiBootTar(out, ctx, cmdline)
	} else {
		return buildGrubBootTar(out, ctx)
	}
}

func buildGrubBootTar(out io.Writer, ctx *renderContext) (err error) {
	arch := tar.NewWriter(out)
	defer arch.Close()

	archAdd := func(path string, ba []byte) (err error) {
		err = arch.WriteHeader(&tar.Header{Name: path, Mode: 0640, Size: int64(len(ba))})
		if err != nil {
			return
		}
		_, err = arch.Write(ba)
		return
	}

	// kernel
	kernelPath, err := distFetch("kernels", ctx.Host.Kernel)
	if err != nil {
		return
	}

	kernelBytes, err := os.ReadFile(kernelPath)
	if err != nil {
		return
	}

	err = archAdd("current/vmlinuz", kernelBytes)
	if err != nil {
		return
	}

	// initrd
	initrd := new(bytes.Buffer)
	err = buildInitrd(initrd, ctx)
	if err != nil {
		return
	}

	err = archAdd("current/initrd", initrd.Bytes())
	if err != nil {
		return
	}

	// done
	return nil
}

func buildUkiBootTar(out io.Writer, ctx *renderContext, cmdline string) (err error) {
	arch := tar.NewWriter(out)
	defer arch.Close()

	archAdd := func(path string, ba []byte) (err error) {
		err = arch.WriteHeader(&tar.Header{Name: path, Mode: 0640, Size: int64(len(ba))})
		if err != nil {
			return
		}
		_, err = arch.Write(ba)
		return
	}

	// UKI
	uki := new(bytes.Buffer)
	err = buildUki(uki, ctx, true, cmdline)
	if err != nil {
		return
	}

	err = archAdd("EFI/BOOT/BOOTX64.EFI", uki.Bytes())
	if err != nil {
		return
	}

	// done
	return nil
}