package main

import (
        "fmt"
        "io"
        "net"
        "time"
        "net/http"
        "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 requireRoot() {
        if os.Geteuid() != 0 {
                fmt.Fprintln(os.Stderr, "mnmivm must be run as root. Try: `sudo mnmivm` or login as the root user")
                os.Exit(1)
        }
}

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("  mnmivm create <name> --os <ubuntu|debian|fedora|alpine> --pubkey-path <path>")
        fmt.Println("  mnmivm start <name>")
        fmt.Println("  mnmivm stop <name>")
	fmt.Println("  mnmivm update-cloud <name> --pubkey-path <path>")
        fmt.Println("  mnmivm delete <name>")
        fmt.Println("  mnmivm 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 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 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 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 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 updateCloudVM(name, pubKeyPath string) {
        dir := vmDir(name)

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

        if vmRunning(name) {
                panic("VM is currently running. Stop it before updating cloud-init.")
        }

        osName := readFileTrim(filepath.Join(dir, "os.name"))
        if osName == "-" {
                panic("cannot determine OS for VM")
        }

        // overwrite stored pubkey
        keyData, err := os.ReadFile(pubKeyPath)
        if err != nil {
                panic(err)
        }

        if err := os.WriteFile(vmPubKeyPath(name), keyData, 0644); err != nil {
                panic(err)
        }

        // regenerate seed ISO
        createCloudInitSeed(name, osName, vmPubKeyPath(name))

        fmt.Println("Cloud-init updated for VM:", name)
        fmt.Println("Changes will apply on next boot.")
}

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

        instanceID := fmt.Sprintf("%s-%d", name, time.Now().Unix())

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

ssh_pwauth: false
disable_root: true

write_files:
  - path: /home/%s/.ssh/authorized_keys
    owner: %s:%s
    permissions: '0600'
    content: |
      %s

runcmd:
  - chown -R %s:%s /home/%s/.ssh
`, img.User, img.User, img.User, img.User, sshKey,
   img.User, img.User, img.User)

        metaData := fmt.Sprintf(
                "instance-id: %s\nlocal-hostname: %s\n",
                instanceID,
                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")

        if err := os.WriteFile(userDataPath, []byte(userData), 0644); err != nil {
                panic(err)
        }
        if err := os.WriteFile(metaDataPath, []byte(metaData), 0644); err != nil {
                panic(err)
        }

        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 renderBoxTable(headers []string, rows [][]string) {
        colWidths := make([]int, len(headers))

        // compute column widths
        for i, h := range headers {
                colWidths[i] = len(h)
        }
        for _, row := range rows {
                for i, cell := range row {
                        if len(cell) > colWidths[i] {
                                colWidths[i] = len(cell)
                        }
                }
        }

        // helpers
        repeat := func(s string, n int) string {
                return strings.Repeat(s, n)
        }

        drawRow := func(left, mid, right string) {
                fmt.Print(left)
                for i, w := range colWidths {
                        fmt.Print(repeat("─", w+2))
                        if i < len(colWidths)-1 {
                                fmt.Print(mid)
                        }
                }
                fmt.Println(right)
        }

        // top border
        drawRow("┌", "┬", "┐")

        // header row
        fmt.Print("│")
        for i, h := range headers {
                fmt.Printf(" %-*s │", colWidths[i], h)
        }
        fmt.Println()

        // header separator
        drawRow("├", "┼", "┤")

        // data rows
        for _, row := range rows {
                fmt.Print("│")
                for i, cell := range row {
                        fmt.Printf(" %-*s │", colWidths[i], cell)
                }
                fmt.Println()
        }

        // bottom border
        drawRow("└", "┴", "┘")
}


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

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

        keyData, _ := os.ReadFile(pubKeyPath)
        os.WriteFile(vmPubKeyPath(name), keyData, 0644)
        os.WriteFile(filepath.Join(vmDir(name), "os.name"), []byte(osName), 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, vmPubKeyPath(name))
        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))
        }

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

        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", sshPort),
                "-device", "virtio-net-pci,netdev=net0",
                "-pidfile", pidFile(name),
                "-daemonize",
        )

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

        os.WriteFile(filepath.Join(vmDir(name), "ssh.port"), []byte(fmt.Sprint(sshPort)), 0644)
        os.WriteFile(filepath.Join(vmDir(name), "vnc.port"), []byte(fmt.Sprint(vncPort)), 0644)

        fmt.Println("VM started:", name)
        fmt.Printf("SSH: ssh %s@localhost -p %d\n", osImages[readFileTrim(filepath.Join(vmDir(name), "os.name"))].User, sshPort)
        fmt.Printf("VNC: vnc://0.0.0.0:%d\n", vncPort)
}

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 deleteVM(name string) {
        dir := vmDir(name)
        if _, err := os.Stat(dir); err != nil {
                panic("VM does not exist")
        }

        if data, err := os.ReadFile(pidFile(name)); err == nil {
                exec.Command("kill", strings.TrimSpace(string(data))).Run()
                os.Remove(pidFile(name))
        }

        fmt.Printf("\nWARNING: Permanently delete VM \"%s\"?\nType YES to confirm: ", name)
        var confirm string
        fmt.Scanln(&confirm)
        if confirm != "YES" {
                fmt.Println("Aborted.")
                return
        }

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

func readFileTrim(path string) string {
        data, err := os.ReadFile(path)
        if err != nil {
                return "-"
        }
        return strings.TrimSpace(string(data))
}

func vmRunning(name string) bool {
        data, err := os.ReadFile(pidFile(name))
        if err != nil {
                return false
        }
        _, err = os.Stat("/proc/" + strings.TrimSpace(string(data)))
        return err == nil
}

func listVMs() {
        entries, _ := os.ReadDir(vmDirBase)

        headers := []string{
                "NAME",
                "STATE",
                "OS",
                "SSH",
                "VNC",
                "PUBKEY",
        }

        rows := [][]string{}

        for _, e := range entries {
                name := e.Name()
                dir := vmDir(name)

                state := "stopped"
                if vmRunning(name) {
                        state = "running"
                }

                osName := readFileTrim(filepath.Join(dir, "os.name"))

                sshPort := readFileTrim(filepath.Join(dir, "ssh.port"))
                vncPort := readFileTrim(filepath.Join(dir, "vnc.port"))

                ssh := "-"
                if sshPort != "-" {
                        ssh = "localhost:" + sshPort
                }

                vnc := "-"
                if vncPort != "-" {
                        vnc = "0.0.0.0:" + vncPort
                }

                rows = append(rows, []string{
                        name,
                        state,
                        osName,
                        ssh,
                        vnc,
                        vmPubKeyPath(name),
                })
        }

        if len(rows) == 0 {
                fmt.Println("No VMs found.")
                return
        }

        renderBoxTable(headers, rows)
}


func main() {
	requireRoot()
        if len(os.Args) < 2 {
                usage()
        }

        ensureDirs()

        switch os.Args[1] {
        case "create":
                createVM(os.Args[2], parseArg(os.Args, "--os"), parseArg(os.Args, "--pubkey-path"))
        case "start":
                startVM(os.Args[2])
        case "stop":
                stopVM(os.Args[2])
        case "delete":
                deleteVM(os.Args[2])
        case "list":
                listVMs()
	case "update-cloud":
		if len(os.Args) < 3 {
			usage()
		}
		updateCloudVM(os.Args[2], parseArg(os.Args, "--pubkey-path"))

        default:
                usage()
        }
}