package main

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

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

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

	fmt.Printf("Creating working 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)
	}

	// 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)
	}

	// Install a kernel using debootstrap or apt
	fmt.Println("Installing kernel and bootloader...")
	
	// Bind mount necessary directories for package installation
	bindDirs := []string{"/dev", "/proc", "/sys"}
	for _, dir := range bindDirs {
		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)
		}
	}

	// Update package lists
	fmt.Println("Updating package lists...")
	cmd = exec.Command("sudo", "chroot", mountPoint, "apt", "update")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Warning: apt update failed: %v\n", err)
	}

	// Install kernel and essential packages
	fmt.Println("Installing kernel and essential packages...")
	kernelPackages := []string{
		"linux-image-amd64",      // Main kernel
		"linux-headers-amd64",    // Kernel headers
		"grub-pc",                // GRUB bootloader
		"grub-pc-bin",            // GRUB binaries
		"os-prober",              // OS detection
		"initramfs-tools",        // Initial RAM disk
		"systemd-sysv",           // Systemd
		"systemd",                 // Systemd
		"udev",                    // Device management
		"console-setup",          // Console setup
		"keyboard-configuration", // Keyboard config
		"locales",                // Locales
		"tzdata",                 // Timezone data
	}

	for _, pkg := range kernelPackages {
		fmt.Printf("Installing %s...\n", pkg)
		cmd = exec.Command("sudo", "chroot", mountPoint, "apt", "install", "-y", "--no-install-recommends", pkg)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		if err := cmd.Run(); err != nil {
			fmt.Printf("Warning: failed to install %s: %v\n", pkg, err)
		}
	}

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

	// Install GRUB bootloader
	fmt.Println("Installing GRUB bootloader...")
	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
	fmt.Println("Generating GRUB configuration...")
	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)
	}

	// Create a simple fstab
	fmt.Println("Creating fstab...")
	fstabContent := `# /etc/fstab for particle-os
/dev/sda1	/	ext4	rw,errors=remount-ro	0	1
proc	/proc	proc	defaults	0	0
sysfs	/sys	sysfs	defaults	0	0
devpts	/dev/pts	devpts	gid=5,mode=620	0	0
tmpfs	/run	tmpfs	defaults	0	0
`
	fstabFile := filepath.Join(mountPoint, "etc", "fstab")
	cmd = exec.Command("sudo", "sh", "-c", fmt.Sprintf("cat > %s << 'EOF'\n%sEOF", fstabFile, fstabContent))
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Warning: could not create fstab: %v\n", err)
	}

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

	fmt.Printf("\n✅ Working bootable image created successfully: %s\n", outputPath)
	fmt.Printf("Image size: %s\n", imageSize)
	fmt.Printf("Filesystem: ext4\n")
	fmt.Printf("Bootloader: GRUB with kernel\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)
}