package main

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
)

const (
	rootDir   = "/var/lib/microvm"
	imageDir  = rootDir + "/images"
	vmDirBase = rootDir + "/vms"

	defaultImgURL  = "https://cloud-images.ubuntu.com/noble/current/noble-server-cloudimg-amd64.img"
	defaultImgName = "ubuntu-24.04-cloudinit.qcow2"

	memMB = "1024"
	cpus  = "1"
)

func ensureDirs() {
	for _, d := range []string{imageDir, vmDirBase} {
		if err := os.MkdirAll(d, 0755); err != nil {
			panic(err)
		}
	}
}

func baseImagePath() string {
	return filepath.Join(imageDir, defaultImgName)
}

func vmDir(name string) string   { return filepath.Join(vmDirBase, name) }
func pidFile(name string) string { return filepath.Join(vmDir(name), "vm.pid") }
func diskPath(name string) string {
	return filepath.Join(vmDir(name), "disk.qcow2")
}
func seedISOPath(name string) string {
	return filepath.Join(vmDir(name), "seed.iso")
}

func sshPort(name string) int {
	if strings.HasPrefix(name, "vm") {
		if n, err := strconv.Atoi(strings.TrimPrefix(name, "vm")); err == nil {
			return 2220 + n
		}
	}
	return 2222
}

func readSSHPublicKey() string {
	home, err := os.UserHomeDir()
	if err != nil {
		panic(err)
	}

	keys := []string{
		filepath.Join(home, ".ssh", "id_ed25519.pub"),
		filepath.Join(home, ".ssh", "id_rsa.pub"),
	}

	for _, k := range keys {
		if data, err := os.ReadFile(k); err == nil {
			return strings.TrimSpace(string(data))
		}
	}

	panic("no SSH public key found in ~/.ssh")
}

func createCloudInitSeed(name string) {
	sshKey := readSSHPublicKey()

	userData := fmt.Sprintf(`#cloud-config
users:
  - name: ubuntu
    sudo: ALL=(ALL) NOPASSWD:ALL
    shell: /bin/bash
    ssh_authorized_keys:
      - %s

ssh_pwauth: false
disable_root: true
`, sshKey)

	metaData := fmt.Sprintf(`instance-id: %s
local-hostname: %s
`, name, name)

	tmpDir, err := os.MkdirTemp("", "cloudinit")
	if err != nil {
		panic(err)
	}
	defer os.RemoveAll(tmpDir)

	userDataPath := filepath.Join(tmpDir, "user-data")
	metaDataPath := filepath.Join(tmpDir, "meta-data")

	os.WriteFile(userDataPath, []byte(userData), 0644)
	os.WriteFile(metaDataPath, []byte(metaData), 0644)

	cmd := exec.Command(
		"genisoimage",
		"-output", seedISOPath(name),
		"-volid", "cidata",
		"-joliet",
		"-rock",
		userDataPath,
		metaDataPath,
	)

	if out, err := cmd.CombinedOutput(); err != nil {
		panic("cloud-init ISO creation failed:\n" + string(out))
	}
}

func ensureBaseImage() {
	if _, err := os.Stat(baseImagePath()); err == nil {
		return
	}

	fmt.Println("Downloading Ubuntu 24.04 cloud image...")

	resp, err := http.Get(defaultImgURL)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	out, err := os.Create(baseImagePath())
	if err != nil {
		panic(err)
	}
	defer out.Close()

	if _, err := io.Copy(out, resp.Body); err != nil {
		panic(err)
	}
}

func createVM(name string) {
	ensureBaseImage()
	os.MkdirAll(vmDir(name), 0755)

	cmd := exec.Command(
		"qemu-img", "create",
		"-f", "qcow2",
		"-F", "qcow2",
		"-b", baseImagePath(),
		diskPath(name),
	)

	if out, err := cmd.CombinedOutput(); err != nil {
		panic(string(out))
	}

	createCloudInitSeed(name)

	fmt.Println("VM created:", name)
}

func startVM(name string) {
	if data, err := os.ReadFile(pidFile(name)); err == nil {
		pid := strings.TrimSpace(string(data))
		if _, err := os.Stat("/proc/" + pid); err == nil {
			panic("VM already running")
		}
		os.Remove(pidFile(name))
	}

	port := sshPort(name)

	cmd := exec.Command(
	    "qemu-system-x86_64",
	    "-enable-kvm",
	    "-machine", "q35",
	    "-cpu", "host",
	    "-m", memMB,
	    "-smp", cpus,

	    // VNC-only console
	    "-vga", "virtio",
	    "-display", "vnc=0.0.0.0:1",

	    // Main disk (cloud image)
	    "-drive", fmt.Sprintf("file=%s,if=virtio,format=qcow2", diskPath(name)),

	    // Cloud-init seed ISO
	    "-drive", fmt.Sprintf("file=%s,if=virtio,format=raw", seedISOPath(name)),

	    // User-mode networking + SSH forward
	    "-netdev", fmt.Sprintf("user,id=net0,hostfwd=tcp::%d-:22", port),
	    "-device", "virtio-net-pci,netdev=net0",

	    // Boot from disk (GRUB inside image)
	    "-boot", "order=c",

	    "-pidfile", pidFile(name),
	    "-daemonize",
	)


	if out, err := cmd.CombinedOutput(); err != nil {
		panic("qemu failed:\n" + string(out))
	}

	fmt.Println("VM started:", name)
	fmt.Printf("VNC: localhost:5901\n")
	fmt.Printf("SSH: ssh ubuntu@localhost -p %d\n", port)
}

func stopVM(name string) {
	data, err := os.ReadFile(pidFile(name))
	if err != nil {
		panic("VM not running")
	}
	pid := strings.TrimSpace(string(data))
	exec.Command("kill", pid).Run()
	os.Remove(pidFile(name))
	fmt.Println("VM stopped:", name)
}

func listVMs() {
	entries, _ := os.ReadDir(vmDirBase)
	for _, e := range entries {
		fmt.Println(e.Name())
	}
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("usage: microvm <create|start|stop|list> <name>")
		os.Exit(1)
	}

	ensureDirs()

	switch os.Args[1] {
	case "create":
		createVM(os.Args[2])
	case "start":
		startVM(os.Args[2])
	case "stop":
		stopVM(os.Args[2])
	case "list":
		listVMs()
	default:
		fmt.Println("usage: microvm <create|start|stop|list> <name>")
	}
}