package main

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

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

	memMB = "1024"
	cpus  = "1"
)

type OSImage struct {
	Name     string
	URL      string
	Filename string
	User     string
}

var osImages = map[string]OSImage{
	"ubuntu": {
		Name:     "ubuntu",
		URL:      "https://cloud-images.ubuntu.com/noble/current/noble-server-cloudimg-amd64.img",
		Filename: "ubuntu-24.04.qcow2",
		User:     "ubuntu",
	},
	"debian": {
		Name:     "debian",
		URL:      "https://cloud.debian.org/images/cloud/trixie/20251117-2299/debian-13-generic-amd64-20251117-2299.qcow2",
		Filename: "debian-13.qcow2",
		User:     "debian",
	},
	"fedora": {
		Name:     "fedora",
		URL:      "https://download.fedoraproject.org/pub/fedora/linux/releases/43/Cloud/x86_64/images/Fedora-Cloud-Base-Generic-43-1.6.x86_64.qcow2",
		Filename: "fedora-43.qcow2",
		User:     "fedora",
	},
	"alpine": {
		Name:     "alpine",
		URL:      "https://dl-cdn.alpinelinux.org/alpine/v3.22/releases/cloud/generic_alpine-3.22.2-x86_64-bios-cloudinit-r0.qcow2",
		Filename: "alpine-3.22.qcow2",
		User:     "alpine",
	},
}

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

func usage() {
	fmt.Println("usage:")
	fmt.Println("  microvm create <name> --os <ubuntu|debian|fedora|alpine> --pubkey-path <path>")
	fmt.Println("  microvm start <name>")
	fmt.Println("  microvm stop <name>")
	fmt.Println("  microvm delete <name>")
	fmt.Println("  microvm list")
	os.Exit(1)
}

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 vmPubKeyPath(name string) string { return filepath.Join(vmDir(name), "pubkey.pub") }

func imagePath(osName string) string {
	return filepath.Join(imageDir, osImages[osName].Filename)
}

func deleteVM(name string) {
        dir := vmDir(name)

        if _, err := os.Stat(dir); err != nil {
                panic("VM does not exist: " + name)
        }

        // If running, stop it first
        if data, err := os.ReadFile(pidFile(name)); err == nil {
                pid := strings.TrimSpace(string(data))
                if _, err := os.Stat("/proc/" + pid); err == nil {
                        fmt.Println("VM is running, stopping it first...")
                        exec.Command("kill", pid).Run()
                }
                os.Remove(pidFile(name))
        }

        fmt.Printf(`
WARNING: This will permanently delete the VM "%s"
Path: %s

Type YES (all caps) to confirm: `, name, dir)

        var confirm string
        fmt.Scanln(&confirm)

        if confirm != "YES" {
                fmt.Println("Aborted. VM not deleted.")
                return
        }

        if err := os.RemoveAll(dir); err != nil {
                panic(err)
        }

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


func findFreePort() int {
        l, err := net.Listen("tcp", "127.0.0.1:0")
        if err != nil {
                panic(err)
        }
        defer l.Close()
        return l.Addr().(*net.TCPAddr).Port
}


func parseArg(args []string, key string) string {
	for i := 0; i < len(args)-1; i++ {
		if args[i] == key {
			return args[i+1]
		}
	}
	panic("missing required argument: " + key)
}

func readPublicKey(path string) string {
	data, err := os.ReadFile(path)
	if err != nil {
		panic(err)
	}
	key := strings.TrimSpace(string(data))
	if !strings.HasPrefix(key, "ssh-") {
		panic("invalid SSH public key")
	}
	return key
}

func ensureBaseImage(osName string) {
	img := osImages[osName]
	path := imagePath(osName)

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

	fmt.Println("Downloading", img.Name, "cloud image...")

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

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

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

func createCloudInitSeed(name, osName, pubKeyPath string) {
	img := osImages[osName]
	sshKey := readPublicKey(pubKeyPath)

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

ssh_pwauth: false
disable_root: true
`, img.User, sshKey)

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

	tmpDir, _ := os.MkdirTemp("", "cloudinit")
	defer os.RemoveAll(tmpDir)

	os.WriteFile(filepath.Join(tmpDir, "user-data"), []byte(userData), 0644)
	os.WriteFile(filepath.Join(tmpDir, "meta-data"), []byte(metaData), 0644)

	cmd := exec.Command(
		"genisoimage",
		"-output", seedISOPath(name),
		"-volid", "cidata",
		"-joliet",
		"-rock",
		filepath.Join(tmpDir, "user-data"),
		filepath.Join(tmpDir, "meta-data"),
	)

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

func createVM(name, osName, pubKeyPath string) {
	if _, ok := osImages[osName]; !ok {
		panic("unsupported OS: " + osName)
	}

	ensureBaseImage(osName)
	os.MkdirAll(vmDir(name), 0755)

	keyDst := vmPubKeyPath(name)
	data, _ := os.ReadFile(pubKeyPath)
	os.WriteFile(keyDst, data, 0644)

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

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

	createCloudInitSeed(name, osName, keyDst)

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

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

	port :=  findFreePort()
	vncPort := findFreePort()
	vncDisplay := vncPort - 5900
	if vncDisplay < 0 {
		panic("invalid VNC port allocation")
	}


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

		"-vga", "virtio",
		"-display", fmt.Sprintf("vnc=0.0.0.0:%d", vncDisplay),

		"-drive", fmt.Sprintf("file=%s,if=virtio,format=qcow2", diskPath(name)),
		"-drive", fmt.Sprintf("file=%s,if=virtio,format=raw", seedISOPath(name)),

		"-netdev", fmt.Sprintf("user,id=net0,hostfwd=tcp::%d-:22", port),
		"-device", "virtio-net-pci,netdev=net0",

		"-boot", "order=c",
		"-pidfile", pidFile(name),
		"-daemonize",
	)

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

	fmt.Println("VM started:", name)
	fmt.Printf("VNC: vnc://127.0.0.1:%d\n", vncPort)
	fmt.Printf("SSH: ssh <user>@localhost -p %d\n", port)
}

func stopVM(name string) {
	data, _ := os.ReadFile(pidFile(name))
	exec.Command("kill", strings.TrimSpace(string(data))).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 {
		usage()
	}

	ensureDirs()

	switch os.Args[1] {
	case "create":
		if len(os.Args) < 3 {
			usage()
		}
		osName := parseArg(os.Args, "--os")
		pubKey := parseArg(os.Args, "--pubkey-path")
		createVM(os.Args[2], osName, pubKey)

	case "start":
		startVM(os.Args[2])

	case "stop":
		stopVM(os.Args[2])

	case "delete":
        if len(os.Args) < 3 {
                usage()
        }
        deleteVM(os.Args[2])

	case "list":
		listVMs()

	default:
		usage()
	}
}