#!/bin/bash

set -e

# Default configuration
HYPERVISOR="${HYPERVISOR:-qemu}" # qemu or ch (cloud-hypervisor)
NUM_DISKS="${NUM_DISKS:-3}"
DISK_SIZE="${DISK_SIZE:-10G}"
DISK_PREFIX="zosvol"
KERNEL="dist/vmlinuz.efi"
MEMORY="2048"
BRIDGE="${BRIDGE:-zosbr}"
TAP_INTERFACE="${TAP_INTERFACE:-zos-tap0}"

# Parse command line arguments
RESET_VOLUMES=false

usage() {
	cat <<EOF
Usage: $0 [OPTIONS]

Options:
    -H, --hypervisor <qemu|ch>  Choose hypervisor (default: qemu)
    -d, --disks <N>             Number of disks to create (default: 3)
    -s, --disk-size <SIZE>      Size of each disk (default: 10G)
    -r, --reset                 Delete and recreate all volumes
    -m, --memory <MB>           Memory in MB (default: 2048)
    -b, --bridge <NAME>         Network bridge name (default: zosbr)
    -h, --help                  Show this help message

Environment variables:
    HYPERVISOR              Same as --hypervisor
    NUM_DISKS               Same as --disks
    DISK_SIZE               Same as --disk-size
    BRIDGE                  Same as --bridge

Examples:
    $0 -H qemu -d 5
    $0 --hypervisor ch --reset
    $0 -d 4 -r
    HYPERVISOR=ch NUM_DISKS=4 $0
EOF
	exit 0
}

# Parse arguments using getopt
TEMP=$(getopt -o 'H:d:s:rm:b:h' --long 'hypervisor:,disks:,disk-size:,reset,memory:,bridge:,help' -n "$0" -- "$@")

if [ $? -ne 0 ]; then
	echo "Error parsing arguments. Try '$0 --help' for more information."
	exit 1
fi

eval set -- "$TEMP"
unset TEMP

while true; do
	case "$1" in
	'-H' | '--hypervisor')
		HYPERVISOR="$2"
		shift 2
		;;
	'-d' | '--disks')
		NUM_DISKS="$2"
		shift 2
		;;
	'-s' | '--disk-size')
		DISK_SIZE="$2"
		shift 2
		;;
	'-r' | '--reset')
		RESET_VOLUMES=true
		shift
		;;
	'-m' | '--memory')
		MEMORY="$2"
		shift 2
		;;
	'-b' | '--bridge')
		BRIDGE="$2"
		shift 2
		;;
	'-h' | '--help')
		usage
		;;
	'--')
		shift
		break
		;;
	*)
		echo "Internal error!"
		exit 1
		;;
	esac
done

# Validate hypervisor choice
if [[ "$HYPERVISOR" != "qemu" && "$HYPERVISOR" != "ch" ]]; then
	echo "Error: Invalid hypervisor '$HYPERVISOR'. Must be 'qemu' or 'ch'"
	exit 1
fi

if [[ "$HYPERVISOR" == "qemu" ]]; then
	DISK_SUFFIX="qcow"
	DISK_FORMAT="qcow2"
else
	DISK_SUFFIX="raw"
	DISK_FORMAT="raw"
fi

# Validate kernel exists
if [[ ! -f "$KERNEL" ]]; then
	echo "Error: Kernel not found at $KERNEL"
	echo "Run the build first: ./scripts/build.sh"
	exit 1
fi

# Setup TAP interface for cloud-hypervisor
setup_tap_interface() {
	local tap="$1"
	local bridge="$2"

	echo "Setting up TAP interface $tap for cloud-hypervisor..."
	if [[ -z "${CH_TAP_IP:-}" && -z "${CH_TAP_MASK:-}" ]]; then
		echo "  Reminder: set CH_TAP_IP and CH_TAP_MASK to configure $tap and silence cloud-hypervisor IP warnings."
	fi

	# Check if bridge exists
	if ! ip link show "$bridge" &>/dev/null; then
		echo "Warning: Bridge $bridge does not exist. Create it with:"
		echo "  sudo ip link add name $bridge type bridge"
		echo "  sudo ip link set $bridge up"
		echo ""
		read -p "Continue anyway? (y/N) " -n 1 -r
		echo
		if [[ ! $REPLY =~ ^[Yy]$ ]]; then
			exit 1
		fi
	fi

	# Check if TAP already exists
	if ip link show "$tap" &>/dev/null; then
		echo "  TAP interface $tap already exists"
		# Check if it's already attached to bridge
		if bridge link show | grep -q "$tap"; then
			echo "  $tap is already attached to bridge"
			return 0
		else
			echo "  Attaching $tap to bridge $bridge"
			sudo ip link set dev "$tap" master "$bridge" 2>/dev/null || true
			return 0
		fi
	fi

	# Create new TAP interface
	echo "  Creating TAP interface $tap"
	sudo ip tuntap add "$tap" mode tap user "$(whoami)"

	# Bring it up
	echo "  Bringing up $tap"
	sudo ip link set "$tap" up

	# Attach to bridge
	echo "  Attaching $tap to bridge $bridge"
	sudo ip link set dev "$tap" master "$bridge"

	echo "  TAP interface setup complete"
}

# Cleanup TAP interface on exit (for cloud-hypervisor)
cleanup_tap_interface() {
	local tap="$1"
	if [[ -n "$tap" ]] && ip link show "$tap" &>/dev/null; then
		echo "Cleaning up TAP interface $tap"
		sudo ip link delete "$tap" 2>/dev/null || true
	fi
}

# Reset volumes if requested
if [[ "$RESET_VOLUMES" == "true" ]]; then
	echo "Resetting volumes..."
	for i in $(seq 0 $((NUM_DISKS - 1))); do
		for suffix in qcow raw; do
			vol="${DISK_PREFIX}${i}.${suffix}"
			if [[ -f "$vol" ]]; then
				echo "  Removing $vol"
				rm -f "$vol"
			fi
		done
	done
fi

# Create disk volumes if they don't exist
echo "Ensuring $NUM_DISKS disk volume(s) exist..."
for i in $(seq 0 $((NUM_DISKS - 1))); do
	vol="${DISK_PREFIX}${i}.${DISK_SUFFIX}"
	if [[ ! -f "$vol" ]]; then
		echo "  Creating $vol (size: $DISK_SIZE, format: $DISK_FORMAT)"
		qemu-img create -f "$DISK_FORMAT" "$vol" "$DISK_SIZE" >/dev/null
	else
		echo "  $vol already exists"
	fi
done

# Build disk arguments based on hypervisor
DISK_ARGS=()
if [[ "$HYPERVISOR" == "qemu" ]]; then
	for i in $(seq 0 $((NUM_DISKS - 1))); do
		vol="${DISK_PREFIX}${i}.${DISK_SUFFIX}"
		DISK_ARGS+=("-drive" "file=$vol,format=$DISK_FORMAT,if=virtio")
	done
elif [[ "$HYPERVISOR" == "ch" ]]; then
	# cloud-hypervisor requires comma-separated disk list in single --disk argument
	# Use raw format (no compression)
	CH_DISKS=""
	for i in $(seq 0 $((NUM_DISKS - 1))); do
		vol="${DISK_PREFIX}${i}.${DISK_SUFFIX}"
		if [[ -z "$CH_DISKS" ]]; then
			CH_DISKS="path=$vol,iommu=on"
		else
			CH_DISKS="${CH_DISKS},path=$vol,iommu=on"
		fi
	done
	if [[ -n "$CH_DISKS" ]]; then
		DISK_ARGS+=("--disk" "$CH_DISKS")
	fi
fi

# Launch the appropriate hypervisor
echo "Launching with $HYPERVISOR hypervisor..."
echo ""

if [[ "$HYPERVISOR" == "qemu" ]]; then
	exec qemu-system-x86_64 \
		-kernel "$KERNEL" \
		-m "$MEMORY" \
		-enable-kvm \
		-cpu host \
		-net nic,model=virtio \
		-net bridge,br="$BRIDGE" \
		"${DISK_ARGS[@]}" \
		-chardev stdio,id=char0,mux=on,logfile=serial.log \
		-serial chardev:char0 \
		-mon chardev=char0 \
		-append "console=ttyS0,115200n8 console=tty"

elif [[ "$HYPERVISOR" == "ch" ]]; then
	# Setup TAP interface for cloud-hypervisor
	setup_tap_interface "$TAP_INTERFACE" "$BRIDGE"

	# Cleanup on exit
	trap "cleanup_tap_interface '$TAP_INTERFACE'" EXIT

	exec cloud-hypervisor \
		--kernel "$KERNEL" \
		--memory size="${MEMORY}M" \
		--cpus boot=2 \
		--net "tap=$TAP_INTERFACE" \
		"${DISK_ARGS[@]}" \
		--console off \
		--serial tty \
		--cmdline "console=ttyS0,115200n8"
fi