package main

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
)

func buildBootImg(out io.Writer, ctx *renderContext, uki bool, cmdline string) (err error) {
	bootImg, err := os.CreateTemp(os.TempDir(), "boot.img-")
	if err != nil {
		return
	}
	defer rmTempFile(bootImg)

	err = setupBootImage(bootImg, ctx, uki, cmdline)
	if err != nil {
		return
	}

	// send the result
	bootImg.Seek(0, io.SeekStart)
	io.Copy(out, bootImg)
	return
}

func buildBootImgGZ(out io.Writer, ctx *renderContext, uki bool, cmdline string) (err error) {
	gzOut := gzip.NewWriter(out)

	if err = buildBootImg(gzOut, ctx, uki, cmdline); err != nil {
		return
	}

	gzOut.Close()
	return
}

func buildBootImgQemuConvert(out io.Writer, ctx *renderContext, format string, uki bool, cmdline string) (err error) {
	imgPath, err := func() (imgPath string, err error) {
		bootImg, err := os.CreateTemp(os.TempDir(), "boot.img-")
		if err != nil {
			return
		}
		defer rmTempFile(bootImg)

		err = setupBootImage(bootImg, ctx, uki, cmdline)
		if err != nil {
			return
		}

		err = bootImg.Sync()
		if err != nil {
			return
		}

		imgPath = bootImg.Name() + "." + format

		err = run("qemu-img", "convert", "-f", "raw", "-O", format, bootImg.Name(), imgPath)
		if err != nil {
			return
		}
		return
	}()

	if imgPath != "" {
		defer os.Remove(imgPath)
	}

	if err != nil {
		return
	}

	// send the result
	img, err := os.Open(imgPath)
	if err != nil {
		return
	}

	io.Copy(out, img)
	return
}
func qemuImgBootImg(format string, uki bool, cmdline string) func(out io.Writer, ctx *renderContext) (err error) {
	return func(out io.Writer, ctx *renderContext) (err error) {
		return buildBootImgQemuConvert(out, ctx, format, uki, cmdline)
	}
}

func setupBootImage(bootImg *os.File, ctx *renderContext, uki bool, cmdline string) (err error) {
	if uki {
		err = baseDiskUki(bootImg, 192)
	} else {
		err = baseDiskGrub(bootImg)
	}
	if err != nil {
		return err
	}

	log.Print("running losetup...")
	cmd := exec.Command("losetup", "--find", "--show", "--partscan", bootImg.Name())
	cmd.Stderr = os.Stderr
	devb, err := cmd.Output()
	if err != nil {
		return
	}

	dev := strings.TrimSpace(string(devb))
	defer func() {
		log.Print("detaching ", dev)
		run("losetup", "-d", dev)
	}()

	log.Print("device: ", dev)
	syncSysToDev()

	tempDir := bootImg.Name() + ".p1.mount"

	err = os.Mkdir(tempDir, 0755)
	if err != nil {
		return
	}

	defer func() {
		log.Print("removing ", tempDir)
		os.RemoveAll(tempDir)
	}()

	devp1 := dev + "p1"

	err = syscall.Mount(devp1, tempDir, "vfat", 0, "")
	if err != nil {
		return fmt.Errorf("failed to mount %s to %s: %v", devp1, tempDir, err)
	}

	defer func() {
		log.Print("unmounting ", tempDir)
		syscall.Unmount(tempDir, 0)
	}()

	// add system elements
	tarOut, tarIn := io.Pipe()
	go func() {
		e := buildBootTar(tarIn, ctx, uki, cmdline)
		tarIn.Close()
		if e != nil {
			err = e
		}
	}()

	defer tarOut.Close()

	tarRd := tar.NewReader(tarOut)

	for {
		hdr, err := tarRd.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		log.Print("tar: extracting ", hdr.Name)

		outPath := filepath.Join(tempDir, hdr.Name)
		os.MkdirAll(filepath.Dir(outPath), 0755)

		f, err := os.Create(outPath)
		if err != nil {
			return err
		}

		_, err = io.Copy(f, tarRd)
		f.Close()

		if err != nil {
			return err
		}
	}

	return
}

func baseDiskUki(image *os.File, sizeMB int16) (err error) {
	err = image.Truncate(0)
	if err != nil {
		return
	}

	size := int64(sizeMB) << 20

	err = image.Truncate(size)
	if err != nil {
		return
	}
	err = image.Sync()
	if err != nil {
		return
	}

	err = run("sgdisk", "--zap", "--new=1:2048:", "--typecode=1:EF00", "--change-name=1:ESP", image.Name())
	if err != nil {
		return
	}

	const (
		espOffset        = 2048 * 512
		gptBackupSectors = 34 // real world data
	)
	espSectors := (size-espOffset)/512 - gptBackupSectors

	cfg := fmt.Sprintf("drive d: file=%q offset=%d\n", image.Name(), espOffset)
	err = mtool(cfg, "mformat", "-F", "-T", fmt.Sprint(espSectors), "d:")

	return
}

var grubSupportVersion = flag.String("grub-support", "1.1.1", "GRUB support version")

func baseDiskGrub(image *os.File) (err error) {
	path, err := distFetch("grub-support", *grubSupportVersion)
	if err != nil {
		return
	}

	baseImage, err := os.Open(path)
	if err != nil {
		return
	}

	defer baseImage.Close()

	baseImageGz, err := gzip.NewReader(baseImage)
	if err != nil {
		return
	}

	defer baseImageGz.Close()
	_, err = io.Copy(image, baseImageGz)
	return
}

func mtool(config string, program string, args ...string) (err error) {
	cmd := exec.Command(program, args...)
	cmd.Env = append(os.Environ(), "MTOOLSRC=/dev/stdin")
	cmd.Stdin = bytes.NewBufferString(config)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func run(program string, args ...string) (err error) {
	cmd := exec.Command(program, args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}
