package main

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

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

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

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

	// Check if this is an OSTree system
	fmt.Println("Checking for OSTree filesystem structure...")
	ostreeDirs := []string{
		filepath.Join(rootfsPath, "usr", "lib", "ostree-boot"),
		filepath.Join(rootfsPath, "boot", "ostree"),
		filepath.Join(rootfsPath, "usr", "lib", "systemd"),
		filepath.Join(rootfsPath, "usr", "bin", "bootc"),
		filepath.Join(rootfsPath, "usr", "bin", "bootupd"),
	}

	isOstree := false
	for _, dir := range ostreeDirs {
		if _, err := os.Stat(dir); err == nil {
			fmt.Printf("Found OSTree component: %s\n", dir)
			isOstree = true
		}
	}

	if !isOstree {
		fmt.Println("Warning: This doesn't appear to be an OSTree system")
		fmt.Println("Falling back to standard bootable image creation...")
	}

	// 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 BIOS Boot Partition (1MB) for GRUB
	fmt.Println("Creating BIOS Boot Partition...")
	cmd = exec.Command("sudo", "parted", loopDevice, "mkpart", "primary", "1MiB", "2MiB")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Error creating BIOS Boot Partition: %v\n", err)
		os.Exit(1)
	}

	// Set BIOS Boot Partition type flag
	fmt.Println("Setting BIOS Boot Partition type...")
	cmd = exec.Command("sudo", "parted", loopDevice, "set", "1", "bios_grub", "on")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Warning: could not set BIOS Boot Partition type: %v\n", err)
	}

	// Create EFI System Partition (501MB) for UEFI boot
	fmt.Println("Creating EFI System Partition...")
	cmd = exec.Command("sudo", "parted", loopDevice, "mkpart", "primary", "fat32", "2MiB", "503MiB")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Error creating EFI System Partition: %v\n", err)
		os.Exit(1)
	}

	// Set EFI System Partition type flag
	fmt.Println("Setting EFI System Partition type...")
	cmd = exec.Command("sudo", "parted", loopDevice, "set", "2", "esp", "on")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Warning: could not set EFI System Partition type: %v\n", err)
	}

	// Create Boot Partition (1GB) for GRUB and kernel
	fmt.Println("Creating Boot Partition...")
	cmd = exec.Command("sudo", "parted", loopDevice, "mkpart", "primary", "ext4", "503MiB", "1527MiB")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Error creating Boot Partition: %v\n", err)
		os.Exit(1)
	}

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

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

	// Format EFI partition with FAT32
	fmt.Printf("Formatting EFI partition with FAT32...\n")
	efiPartitionDevice := loopDevice + "p2"
	cmd = exec.Command("sudo", "mkfs.fat", "-F", "32", efiPartitionDevice)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Error formatting EFI partition: %v\n", err)
		os.Exit(1)
	}

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

	// Format Root partition with ext4
	fmt.Printf("Formatting Root partition with ext4...\n")
	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 Root 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 Root partition
	fmt.Printf("Mounting Root 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 Root partition: %v\n", err)
		os.Exit(1)
	}

	// Create boot directory first
	fmt.Printf("Creating boot directory structure...\n")
	bootMountPoint := filepath.Join(mountPoint, "boot")
	
	// Create boot directory and set ownership
	cmd = exec.Command("sudo", "mkdir", "-p", bootMountPoint)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Warning: could not create boot directory: %v\n", err)
	}

	// Set ownership to root
	cmd = exec.Command("sudo", "chown", "root:root", bootMountPoint)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Warning: could not set ownership: %v\n", err)
	}

	// Mount Boot partition
	fmt.Printf("Mounting Boot partition to %s...\n", bootMountPoint)
	cmd = exec.Command("sudo", "mount", bootPartitionDevice, bootMountPoint)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Warning: could not mount Boot partition: %v\n", err)
	}

	// Now create EFI directory inside the mounted boot partition
	efiMountPoint := filepath.Join(mountPoint, "boot", "efi")
	cmd = exec.Command("sudo", "mkdir", "-p", efiMountPoint)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Warning: could not create EFI directory: %v\n", err)
	}

	// Set ownership to root
	cmd = exec.Command("sudo", "chown", "root:root", efiMountPoint)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Warning: could not set ownership: %v\n", err)
	}

	// Mount EFI partition
	fmt.Printf("Mounting EFI partition to %s...\n", efiMountPoint)
	cmd = exec.Command("sudo", "mount", efiPartitionDevice, efiMountPoint)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Printf("Warning: could not mount EFI partition: %v\n", err)
	}

	// Clean up mount on exit
	defer func() {
		fmt.Printf("Unmounting all partitions...\n")
		// Unmount in reverse order: EFI, Boot, then Root
		exec.Command("sudo", "umount", efiMountPoint).Run()
		exec.Command("sudo", "umount", bootMountPoint).Run()
		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)
	}

	if isOstree {
		// Handle OSTree-specific setup
		fmt.Println("Setting up OSTree-specific boot configuration...")
		
		// Create OSTree boot directory structure
		ostreeBootDir := filepath.Join(mountPoint, "boot", "ostree")
		if err := os.MkdirAll(ostreeBootDir, 0755); err != nil {
			fmt.Printf("Warning: could not create ostree boot directory: %v\n", err)
		}

		// Check for kernel in OSTree location
		ostreeKernelDir := filepath.Join(mountPoint, "usr", "lib", "ostree-boot")
		if _, err := os.Stat(ostreeKernelDir); err == nil {
			fmt.Printf("Found OSTree kernel directory: %s\n", ostreeKernelDir)
			
			// List what's in the OSTree boot directory
			if files, err := os.ReadDir(ostreeKernelDir); err == nil {
				fmt.Println("OSTree boot contents:")
				for _, file := range files {
					fmt.Printf("  %s\n", file.Name())
				}
			}
		}

		// Create OSTree-specific fstab
		fmt.Println("Creating OSTree fstab...")
		fstabContent := `# /etc/fstab for OSTree particle-os
# Root filesystem (4th partition)
/dev/sda4	/	ext4	rw,errors=remount-ro	0	1
# Boot partition (3rd partition)
/dev/sda3	/boot	ext4	rw,errors=remount-ro	0	2
# EFI System Partition (2nd partition)
/dev/sda2	/boot/efi	vfat	umask=0077,shortname=winnt	0	2
# BIOS Boot Partition (1st partition) - not mounted
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)
		}

		// Try to install GRUB for OSTree
		fmt.Println("Installing GRUB for OSTree...")
		
		// Install GRUB using podman run to access the Debian container with GRUB tools
		// Mount the host's /dev directory to access loop devices
		// Use the actual loop device name that was created
		grubInstallCmd := fmt.Sprintf("grub-install --target=i386-pc --boot-directory=/mnt/boot --root-directory=/mnt %s", loopDevice)
		cmd = exec.Command("sudo", "podman", "run", "--rm", "--privileged", "--tty", 
			"-v", "/dev:/dev:z", 
			"-v", mountPoint+":/mnt:z", 
			"238208cf481a", 
			"bash", "-c", grubInstallCmd)
		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 using podman run
		grubMkconfigCmd := fmt.Sprintf("grub-mkconfig -o /mnt/boot/grub/grub.cfg")
		cmd = exec.Command("sudo", "podman", "run", "--rm", "--privileged", "--tty", 
			"-v", mountPoint+":/mnt:z", 
			"238208cf481a", 
			"bash", "-c", grubMkconfigCmd)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		if err := cmd.Run(); err != nil {
			fmt.Printf("Warning: GRUB config generation failed: %v\n", err)
		} else {
			fmt.Println("GRUB config generated successfully")
		}

		// GRUB installation complete

	} else {
		// Handle standard bootable image setup
		fmt.Println("Setting up standard bootable image...")
		
		// 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 fstab
		fmt.Println("Creating standard 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)
		}

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

	fmt.Printf("\n✅ OSTree-aware bootable image created successfully: %s\n", outputPath)
	fmt.Printf("Image size: %s\n", imageSize)
	fmt.Printf("Filesystem: ext4\n")
	if isOstree {
		fmt.Printf("Type: OSTree system with GRUB\n")
	} else {
		fmt.Printf("Type: Standard system with extlinux\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)
	
	if isOstree {
		fmt.Printf("\nNote: This is an OSTree system. It should boot using the OSTree boot mechanism.\n")
		fmt.Printf("If it doesn't boot, you may need to install a kernel in the OSTree structure.\n")
	} else {
		fmt.Printf("\nNote: This is a standard system. It should boot using the installed bootloader.\n")
	}
}