package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage: go run fix-bootable-image.go <rootfs-path> [output-path]")
		os.Exit(1)
	}

	rootfsPath := os.Args[1]
	outputPath := "debian-bootable.img"
	if len(os.Args) > 2 {
		outputPath = os.Args[2]
	}

	fmt.Printf("Creating bootable image from rootfs: %s\n", rootfsPath)
	fmt.Printf("Output: %s\n", outputPath)

	// Check if rootfs exists
	if _, err := os.Stat(rootfsPath); os.IsNotExist(err) {
		fmt.Printf("Error: rootfs path does not exist: %s\n", rootfsPath)
		os.Exit(1)
	}

	// Create a 5GB raw disk image
	imageSize := "5G"
	fmt.Printf("Creating %s raw disk image...\n", imageSize)
	
	// Use qemu-img to create the image
	cmd := exec.Command("qemu-img", "create", "-f", "raw", outputPath, imageSize)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	
	if err := cmd.Run(); err != nil {
		fmt.Printf("Error creating image: %v\n", err)
		os.Exit(1)
	}

	// Create a loop device
	fmt.Println("Setting up loop device...")
	cmd = exec.Command("sudo", "losetup", "--find", "--show", outputPath)
	output, err := cmd.Output()
	if err != nil {
		fmt.Printf("Error setting up loop device: %v\n", err)
		os.Exit(1)
	}
	loopDevice := strings.TrimSpace(string(output))
	fmt.Printf("Loop device: %s\n", loopDevice)

	// Clean up loop device on exit
	defer func() {
		fmt.Printf("Cleaning up loop device: %s\n", loopDevice)
		exec.Command("sudo", "losetup", "-d", loopDevice).Run()
	}()

	// Create partition table (GPT)
	fmt.Println("Creating GPT partition table...")
	cmd = exec.Command("sudo", "parted", loopDevice, "mklabel", "gpt")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Error creating partition table: %v\n", err)
		os.Exit(1)
	}

	// Create a single partition
	fmt.Println("Creating partition...")
	cmd = exec.Command("sudo", "parted", loopDevice, "mkpart", "primary", "ext4", "1MiB", "100%")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Error creating partition: %v\n", err)
		os.Exit(1)
	}

	// Get the partition device
	partitionDevice := loopDevice + "p1"
	fmt.Printf("Partition device: %s\n", partitionDevice)

	// Format the partition with ext4
	fmt.Println("Formatting partition with ext4...")
	cmd = exec.Command("sudo", "mkfs.ext4", partitionDevice)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Error formatting partition: %v\n", err)
		os.Exit(1)
	}

	// Create mount point
	mountPoint := "/tmp/particle-os-mount"
	if err := os.MkdirAll(mountPoint, 0755); err != nil {
		fmt.Printf("Error creating mount point: %v\n", err)
		os.Exit(1)
	}

	// Mount the partition
	fmt.Printf("Mounting partition to %s...\n", mountPoint)
	cmd = exec.Command("sudo", "mount", partitionDevice, mountPoint)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Error mounting partition: %v\n", err)
		os.Exit(1)
	}

	// Clean up mount on exit
	defer func() {
		fmt.Printf("Unmounting %s...\n", mountPoint)
		exec.Command("sudo", "umount", mountPoint).Run()
	}()

	// Copy rootfs content
	fmt.Printf("Copying rootfs content from %s to %s...\n", rootfsPath, mountPoint)
	cmd = exec.Command("sudo", "cp", "-a", rootfsPath+"/.", mountPoint+"/")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Error copying rootfs: %v\n", err)
		os.Exit(1)
	}

	// Install GRUB bootloader
	fmt.Println("Installing GRUB bootloader...")
	
	// Bind mount necessary directories for GRUB
	grubDirs := []string{"/dev", "/proc", "/sys"}
	for _, dir := range grubDirs {
		bindMount := filepath.Join(mountPoint, dir)
		if err := exec.Command("sudo", "mount", "--bind", dir, bindMount).Run(); err != nil {
			fmt.Printf("Warning: could not bind mount %s: %v\n", dir, err)
		}
	}

	// Install GRUB
	cmd = exec.Command("sudo", "chroot", mountPoint, "grub-install", "--target=i386-pc", loopDevice)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Warning: GRUB installation failed: %v\n", err)
		fmt.Println("This is expected if GRUB is not available in the rootfs")
	}

	// Generate GRUB config
	cmd = exec.Command("sudo", "chroot", mountPoint, "update-grub")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Warning: GRUB config generation failed: %v\n", err)
	}

	// Unbind mount directories
	for _, dir := range grubDirs {
		bindMount := filepath.Join(mountPoint, dir)
		exec.Command("sudo", "umount", bindMount).Run()
	}

	fmt.Printf("\n✅ Bootable image created successfully: %s\n", outputPath)
	fmt.Printf("Image size: %s\n", imageSize)
	fmt.Printf("Filesystem: ext4\n")
	fmt.Printf("Bootloader: GRUB (if available)\n")
	fmt.Printf("\nTo test the image:\n")
	fmt.Printf("qemu-system-x86_64 -m 2G -drive file=%s,format=raw -nographic\n", outputPath)
}