package main

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

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

	memMB = "1024"
	cpus  = "1"

	baseDiskSize = "12G"

	// Server Edition networking
	bridgeName = "br0"
	lanCIDR    = "192.168.86.0/24"
	lanGW      = "192.168.86.1"
	lanDNS1    = "192.168.86.1"
	lanDNS2    = "8.8.8.8"
)

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",
	},
	"debian-forky": {
		Name:     "debian-forky",
		URL:      "https://cloud.debian.org/images/cloud/forky/daily/20251218-2330/debian-14-generic-amd64-daily-20251218-2330.qcow2",
		Filename: "debian-14.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-se must be run as root. Try: `sudo mnmivm-se` 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-se create <name> --os <ubuntu|debian|debian-forky|fedora|alpine> --pubkey-path <path> --ip <192.168.86.X>")
	fmt.Println("  mnmivm-se start <name>")
	fmt.Println("  mnmivm-se stop <name>")
	fmt.Println("  mnmivm-se update-cloud <name> --pubkey-path <path> [--ip <192.168.86.X>]")
	fmt.Println("  mnmivm-se delete <name>")
	fmt.Println("  mnmivm-se 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 vmIPPath(name string) string     { return filepath.Join(vmDir(name), "vm.ip") }
func vmMACPath(name string) string    { return filepath.Join(vmDir(name), "vm.mac") }

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 parseArgOptional(args []string, key string) string {
	for i := 0; i < len(args)-1; i++ {
		if args[i] == key {
			return args[i+1]
		}
	}
	return ""
}

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 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 ensureBridgeExists() {
	if _, err := os.Stat(filepath.Join("/sys/class/net", bridgeName)); err != nil {
		panic("bridge not found: " + bridgeName + " (expected at /sys/class/net/" + bridgeName + ")")
	}
}

func validateStaticIPv4(ipStr string) {
	ip := net.ParseIP(strings.TrimSpace(ipStr))
	if ip == nil {
		panic("invalid IP address: " + ipStr)
	}

	ip4 := ip.To4()
	if ip4 == nil {
		panic("invalid IPv4 address: " + ipStr)
	}

	_, cidrNet, err := net.ParseCIDR(lanCIDR)
	if err != nil {
		panic("invalid lanCIDR constant: " + lanCIDR)
	}

	if !cidrNet.Contains(ip4) {
		panic("IP must be within " + lanCIDR + ": " + ipStr)
	}

	// Prevent network + broadcast addresses (simple /24-safe guard)
	if ip4[len(ip4)-1] == 0 || ip4[len(ip4)-1] == 255 {
		panic("IP cannot be network or broadcast address: " + ipStr)
	}
}


func generateLocallyAdministeredMAC() string {
	// Locally administered, unicast MAC: set bit1 of first octet, clear multicast bit.
	b := make([]byte, 6)
	_, err := io.ReadFull(strings.NewReader(randomHexBytes(12)), b) // placeholder; replaced below
	if err == nil {
		// not used, fallthrough
	}

	// Real randomness from crypto/rand via /dev/urandom without importing crypto/rand:
	f, err := os.Open("/dev/urandom")
	if err != nil {
		panic(err)
	}
	defer f.Close()
	if _, err := io.ReadFull(f, b); err != nil {
		panic(err)
	}

	b[0] = (b[0] | 0x02) & 0xFE
	return fmt.Sprintf("%02x:%02x:%02x:%02x:%02x:%02x", b[0], b[1], b[2], b[3], b[4], b[5])
}

// Minimal helper to avoid unused import tricks; returns hex chars length n*2 (not cryptographically used).
func randomHexBytes(n int) string {
	const hexd = "0123456789abcdef"
	out := make([]byte, n)
	for i := range out {
		out[i] = hexd[(i*7+3)%len(hexd)]
	}
	return string(out)
}

func updateCloudVM(name, pubKeyPath string, ipOptional 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)
	}

	// optionally overwrite stored IP
	if ipOptional != "" {
		validateStaticIPv4(ipOptional)
		if err := os.WriteFile(vmIPPath(name), []byte(strings.TrimSpace(ipOptional)), 0644); err != nil {
			panic(err)
		}
	}

	// regenerate seed ISO
	ip := readFileTrim(vmIPPath(name))
	mac := readFileTrim(vmMACPath(name))
	if ip == "-" {
		panic("missing VM IP; update-cloud requires --ip or a previously stored vm.ip")
	}
	if mac == "-" {
		panic("missing VM MAC; VM seems corrupted (vm.mac not found)")
	}

	createCloudInitSeed(name, osName, vmPubKeyPath(name), ip, mac)

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

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

	validateStaticIPv4(staticIP)

	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
    lock_passwd: true
    create_home: true

ssh_pwauth: false
disable_root: true

runcmd:
  - mkdir -p /home/%s/.ssh
  - echo '%s' > /home/%s/.ssh/authorized_keys
  - chown -R %s:%s /home/%s/.ssh
  - chmod 700 /home/%s/.ssh
  - chmod 600 /home/%s/.ssh/authorized_keys
`, img.User,
		img.User,
		sshKey,
		img.User,
		img.User, img.User,
		img.User,
		img.User,
		img.User)

	metaData := fmt.Sprintf(
		"instance-id: %s\nlocal-hostname: %s\n",
		instanceID,
		name,
	)

	// Cloud-init network config v2 (netplan-style) with MAC match
	networkConfig := fmt.Sprintf(`version: 2
ethernets:
  eth0:
    match:
      macaddress: %s
    set-name: eth0
    dhcp4: false
    addresses:
      - %s/24
    routes:
      - to: 0.0.0.0/0
        via: %s
    nameservers:
      addresses:
        - %s
        - %s
`, strings.ToLower(strings.TrimSpace(macAddr)),
		strings.TrimSpace(staticIP),
		lanGW,
		lanDNS1,
		lanDNS2,
	)

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

	if err := os.WriteFile(filepath.Join(tmpDir, "user-data"), []byte(userData), 0644); err != nil {
		panic(err)
	}
	if err := os.WriteFile(filepath.Join(tmpDir, "meta-data"), []byte(metaData), 0644); err != nil {
		panic(err)
	}
	if err := os.WriteFile(filepath.Join(tmpDir, "network-config"), []byte(networkConfig), 0644); err != nil {
		panic(err)
	}

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

	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))

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

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

	drawRow("┌", "┬", "┐")

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

	drawRow("├", "┼", "┤")

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

	drawRow("└", "┴", "┘")
}

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

	ensureBaseImage(osName)
	if err := os.MkdirAll(vmDir(name), 0755); err != nil {
		panic(err)
	}

	// Store pubkey + OS
	keyData, err := os.ReadFile(pubKeyPath)
	if err != nil {
		panic(err)
	}
	if err := os.WriteFile(vmPubKeyPath(name), keyData, 0644); err != nil {
		panic(err)
	}
	if err := os.WriteFile(filepath.Join(vmDir(name), "os.name"), []byte(osName), 0644); err != nil {
		panic(err)
	}

	// Store static IP (server edition)
	if err := os.WriteFile(vmIPPath(name), []byte(strings.TrimSpace(staticIP)), 0644); err != nil {
		panic(err)
	}

	// Create and store a stable MAC for matching in cloud-init
	mac := generateLocallyAdministeredMAC()
	if err := os.WriteFile(vmMACPath(name), []byte(mac), 0644); err != nil {
		panic(err)
	}

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

	resizeCmd := exec.Command("qemu-img", "resize", diskPath(name), baseDiskSize)
	if out, err := resizeCmd.CombinedOutput(); err != nil {
		panic("disk resize failed:\n" + string(out))
	}

	// Create cloud-init seed with static network config
	createCloudInitSeed(name, osName, vmPubKeyPath(name), staticIP, mac)

	fmt.Println("VM created:", name, "OS:", osName, "IP:", staticIP, "BRIDGE:", bridgeName)
}

func startVM(name string) {
	ensureBridgeExists()

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

	// VNC stays (console access)
	vncPort := findFreePort()
	vncDisplay := vncPort - 5900
	if vncDisplay < 0 {
		panic("invalid VNC port")
	}

	osName := readFileTrim(filepath.Join(vmDir(name), "os.name"))
	if osName == "-" {
		panic("cannot determine OS for VM")
	}
	ip := readFileTrim(vmIPPath(name))
	if ip == "-" {
		panic("missing vm.ip (static IP not set)")
	}
	mac := readFileTrim(vmMACPath(name))
	if mac == "-" {
		panic("missing vm.mac (MAC not set)")
	}

	// Bridge networking (no user-mode, no hostfwd)
	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)),

		// Bridge-backed NIC on br0, with stable MAC to match cloud-init
		"-netdev", fmt.Sprintf("bridge,id=net0,br=%s", bridgeName),
		"-device", fmt.Sprintf("virtio-net-pci,netdev=net0,mac=%s", strings.TrimSpace(mac)),

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

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

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

	fmt.Println("VM started:", name)
	fmt.Printf("SSH: ssh %s@%s\n", osImages[osName].User, ip)
	fmt.Printf("VNC: vnc://0.0.0.0:%d\n", vncPort)
	fmt.Printf("NET: bridge=%s mac=%s ip=%s\n", bridgeName, strings.TrimSpace(mac), ip)
}

func stopVM(name string) {
	data, err := os.ReadFile(pidFile(name))
	if err != nil {
		panic("VM not running (pidfile missing)")
	}
	_ = 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 listVMs() {
	entries, _ := os.ReadDir(vmDirBase)

	headers := []string{
		"NAME",
		"STATE",
		"OS",
		"IP",
		"VNC",
		"BRIDGE",
	}

	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"))
		ip := readFileTrim(filepath.Join(dir, "vm.ip"))
		vncPort := readFileTrim(filepath.Join(dir, "vnc.port"))

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

		rows = append(rows, []string{
			name,
			state,
			osName,
			ip,
			vnc,
			bridgeName,
		})
	}

	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":
		if len(os.Args) < 3 {
			usage()
		}
		name := os.Args[2]
		osName := parseArg(os.Args, "--os")
		pub := parseArg(os.Args, "--pubkey-path")
		ip := parseArg(os.Args, "--ip")

		createVM(name, osName, pub, ip)

		// Keep behavior: after create, regenerate seed (ensures consistent ISO)
		updateCloudVM(name, pub, ip)

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

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

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

	case "list":
		listVMs()

	case "update-cloud":
		if len(os.Args) < 3 {
			usage()
		}
		name := os.Args[2]
		pub := parseArg(os.Args, "--pubkey-path")
		ip := parseArgOptional(os.Args, "--ip")
		updateCloudVM(name, pub, ip)

	default:
		usage()
	}
}