package main

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

func customizeImage(config *Config) error {
	fmt.Println("Customizing image...")

	if err := checkDependencies(); err != nil {
		return err
	}

	imagePath := config.ImageURL
	if strings.HasSuffix(imagePath, ".qcow2") {
		if err := resizeImage(imagePath, config.DiskSize); err != nil {
			return err
		}
	}

	if config.SSHKey != "" {
		if err := injectSSHKey(imagePath, config.SSHKey); err != nil {
			return err
		}
	}

	if config.GuestAgent {
		if err := installGuestAgent(imagePath); err != nil {
			return err
		}
	}

	fmt.Println("Verifying customized image...")
	if err := verifyImage(imagePath); err != nil {
		return fmt.Errorf("customized image verification failed: %w", err)
	}
	fmt.Println("Customized image verification passed!")

	return nil
}

func checkDependencies() error {
	deps := []string{"qemu-img", "virt-customize"}
	for _, dep := range deps {
		if _, err := exec.LookPath(dep); err != nil {
			return fmt.Errorf("required dependency '%s' not found. Please install libguestfs-tools and qemu-utils", dep)
		}
	}
	return nil
}

func resizeImage(imagePath, size string) error {
	fmt.Printf("Resizing image to %s...\n", size)

	cmd := exec.Command("qemu-img", "resize", imagePath, size)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	if err := cmd.Run(); err != nil {
		return fmt.Errorf("failed to resize image: %w", err)
	}

	return nil
}

func injectSSHKey(imagePath, sshKeyPath string) error {
	fmt.Printf("Injecting SSH key from %s...\n", sshKeyPath)

	keyData, err := os.ReadFile(sshKeyPath)
	if err != nil {
		return fmt.Errorf("failed to read SSH key: %w", err)
	}

	tmpDir := "/tmp/cloud-init-" + filepath.Base(imagePath)
	os.MkdirAll(tmpDir, 0755)
	defer os.RemoveAll(tmpDir)

	keyFile := filepath.Join(tmpDir, "authorized_keys")
	if err := os.WriteFile(keyFile, keyData, 0644); err != nil {
		return fmt.Errorf("failed to write SSH key: %w", err)
	}

	cmd := exec.Command("virt-customize",
		"-a", imagePath,
		"--mkdir", "/root/.ssh",
		"--upload", keyFile+":/root/.ssh/authorized_keys",
		"--chmod", "0600:/root/.ssh/authorized_keys",
	)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	if err := cmd.Run(); err != nil {
		return fmt.Errorf("failed to inject SSH key: %w", err)
	}

	return nil
}

func installGuestAgent(imagePath string) error {
	fmt.Println("Installing QEMU Guest Agent...")

	cmd := exec.Command("virt-customize",
		"-a", imagePath,
		"--install", "qemu-guest-agent",
		"--run-command", "systemctl enable qemu-guest-agent",
	)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	if err := cmd.Run(); err != nil {
		return fmt.Errorf("failed to install qemu-guest-agent: %w", err)
	}

	fmt.Println("QEMU Guest Agent installed and enabled!")
	return nil
}