package main

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

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

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

	fmt.Printf("Creating bootable image with proper permissions 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)
	}

	// Fix permissions after copy
	fmt.Println("Fixing permissions...")
	cmd = exec.Command("sudo", "chown", "-R", "root:root", mountPoint)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Warning: could not fix ownership: %v\n", err)
	}

	// Create minimal bootloader setup
	fmt.Println("Setting up minimal bootloader...")
	
	// Create /boot directory if it doesn't exist
	bootDir := filepath.Join(mountPoint, "boot")
	if err := os.MkdirAll(bootDir, 0755); err != nil {
		fmt.Printf("Warning: could not create boot directory: %v\n", err)
	}

	// Create a simple kernel command line
	cmdline := "root=/dev/sda1 rw console=ttyS0"
	cmdlineFile := filepath.Join(bootDir, "cmdline.txt")
	if err := os.WriteFile(cmdlineFile, []byte(cmdline), 0644); err != nil {
		fmt.Printf("Warning: could not create cmdline.txt: %v\n", err)
	}

	// Create a simple fstab
	fstabContent := fmt.Sprintf("# /etc/fstab for particle-os\n/dev/sda1\t/\text4\trw,errors=remount-ro\t0\t1\n")
	fstabFile := filepath.Join(mountPoint, "etc", "fstab")
	if err := os.WriteFile(fstabFile, []byte(fstabContent), 0644); err != nil {
		fmt.Printf("Warning: could not create fstab: %v\n", err)
	}

	// Install extlinux bootloader
	fmt.Println("Installing extlinux bootloader...")
	
	// Check if extlinux is available
	if _, err := exec.LookPath("extlinux"); err == nil {
		// Install extlinux
		cmd = exec.Command("sudo", "extlinux", "--install", bootDir)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		if err := cmd.Run(); err != nil {
			fmt.Printf("Warning: extlinux installation failed: %v\n", err)
		} else {
			fmt.Println("extlinux installed successfully")
		}
	} else {
		fmt.Println("extlinux not available, skipping bootloader installation")
	}

	// Create a simple syslinux config for direct boot
	syslinuxConfig := fmt.Sprintf(`DEFAULT linux
TIMEOUT 50
PROMPT 0

LABEL linux
  KERNEL /boot/vmlinuz
  APPEND root=/dev/sda1 rw console=ttyS0
`)
	syslinuxFile := filepath.Join(bootDir, "syslinux.cfg")
	if err := os.WriteFile(syslinuxFile, []byte(syslinuxConfig), 0644); err != nil {
		fmt.Printf("Warning: could not create syslinux.cfg: %v\n", err)
	}

	// Install syslinux if available
	if _, err := exec.LookPath("syslinux"); err == nil {
		fmt.Println("Installing syslinux...")
		cmd = exec.Command("sudo", "syslinux", "--install", partitionDevice)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		if err := cmd.Run(); err != nil {
			fmt.Printf("Warning: syslinux installation failed: %v\n", err)
		} else {
			fmt.Println("syslinux installed successfully")
		}
	} else {
		fmt.Println("syslinux not available")
	}

	// Try to install GRUB if available
	if _, err := exec.LookPath("grub-install"); err == nil {
		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)
		} else {
			fmt.Println("GRUB installed successfully")
		}

		// 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()
		}
	} else {
		fmt.Println("GRUB not available")
	}

	fmt.Printf("\n✅ Bootable image with permissions created successfully: %s\n", outputPath)
	fmt.Printf("Image size: %s\n", imageSize)
	fmt.Printf("Filesystem: ext4\n")
	fmt.Printf("Bootloader: extlinux/syslinux/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 -serial stdio\n", outputPath)
}