# vmix CLI — build, copy, and run Windows images
{ pkgs, self, system }:
pkgs.writeShellScriptBin "vmix" ''
  set -euo pipefail

  usage() {
    echo "Usage:"
    echo "  vmix build --image <path> [--generalize key=val,...] [--out-link PATH]"
    echo "  vmix copy  --image <path> [--generalize key=val,...] --to-disk /dev/sdX"
    echo "  vmix copy  --image <path> [--generalize key=val,...] --to-remote-disk user@host:/dev/sdX"
    echo "  vmix run   <qcow2-file> [--mem 4096] [--smp 4] [--ahci]"
    echo ""
    echo "Commands:"
    echo "  build   Build a vmix image (optionally generalized)"
    echo "  copy    Build a vmix image and write it to a disk"
    echo "  run     Boot a qcow2 image with QEMU (SDL if DISPLAY available)"
    echo ""
    echo "Options:"
    echo "  --image PATH              Image path in vmixLib (e.g. windows.images.win10.laptop)"
    echo "  --generalize KEY=VAL,...   Finalize image with comma-separated options:"
    echo "                              username=User  password=  hostname=PC"
    echo "                              timezone=UTC   bgColor=8e8cd8"
    echo "                              delay-oobe-run=true  (OOBE + activation on real hardware)"
    echo "  --to-disk DEVICE          Write to local disk and expand partitions"
    echo "  --to-remote-disk SSH:DEV  Stream to remote disk via SSH and expand partitions"
    echo "                              e.g. root@10.10.10.100:/dev/sda"
    echo "  --ahci                    Use AHCI storage for vmix run (for laptop images)"
    echo "  -y, --yes                 Skip disk write confirmation"
    echo "  --out-link PATH           Symlink for the build result (default: ./result)"
    echo ""
    echo "Examples:"
    echo "  vmix build --image windows.images.win10.laptop \\"
    echo "       --generalize username=Sagar,password=secret,hostname=LAPTOP"
    echo ""
    echo "  vmix copy --image windows.images.win10.laptop \\"
    echo "       --generalize username=Sagar,password=secret,hostname=LAPTOP \\"
    echo "       --to-disk /dev/sda"
    echo ""
    echo "  vmix copy --image windows.images.win10.laptop \\"
    echo "       --generalize username=Sagar,password=secret,hostname=LAPTOP \\"
    echo "       --to-remote-disk root@10.10.10.100:/dev/nvme0n1"
    echo ""
    echo "  vmix run ./result --ahci"
    exit 1
  }

  [[ ''${#} -eq 0 ]] && usage

  COMMAND="$1"; shift

  case "$COMMAND" in
    build|copy|run) ;;
    --help|-h)  usage ;;
    *)          echo "Unknown command: $COMMAND"; usage ;;
  esac

  # --- run command ---
  if [[ "$COMMAND" == "run" ]]; then
    [[ ''${#} -lt 1 ]] && { echo "Error: vmix run <qcow2-file> [--mem 4096] [--smp 4] [--ahci]"; exit 1; }
    RUN_INPUT="$1"; shift
    RUN_MEM=4096
    RUN_SMP=4
    RUN_AHCI=false
    while [[ ''${#} -gt 0 ]]; do
      case "$1" in
        --mem)  RUN_MEM="$2"; shift 2 ;;
        --smp)  RUN_SMP="$2"; shift 2 ;;
        --ahci) RUN_AHCI=true; shift ;;
        *)      echo "Unknown option: $1"; exit 1 ;;
      esac
    done

    RUN_IMAGE="$RUN_INPUT"
    [[ ! -f "$RUN_IMAGE" ]] && { echo "Error: file not found: $RUN_IMAGE"; exit 1; }

    VMIX_DISPLAY="-nographic"
    if [[ -n "''${DISPLAY:-}" ]]; then
      VMIX_DISPLAY="-display sdl"
    fi

    cp ${pkgs.OVMF.fd}/FV/OVMF_VARS.fd /tmp/vmix-run-vars-$$.fd
    chmod +w /tmp/vmix-run-vars-$$.fd
    trap 'rm -f /tmp/vmix-run-vars-$$.fd' EXIT

    echo "=== vmix run ==="
    echo "Image:   $RUN_IMAGE"
    echo "Memory:  $RUN_MEM MB"
    echo "CPUs:    $RUN_SMP"
    echo "Display: $VMIX_DISPLAY"
    echo "AHCI:    $RUN_AHCI"
    echo ""

    exec ${pkgs.qemu}/bin/qemu-system-x86_64 \
      $VMIX_DISPLAY \
      -accel kvm \
      -m "$RUN_MEM" \
      -smp "$RUN_SMP" \
      -cpu host \
      -machine type=q35 \
      -drive if=pflash,format=raw,readonly=on,file=${pkgs.OVMF.fd}/FV/OVMF_CODE.fd \
      -drive if=pflash,format=raw,file=/tmp/vmix-run-vars-$$.fd \
      -rtc base=localtime,clock=host \
      -device qemu-xhci -device usb-tablet \
      $(if [[ "$RUN_AHCI" == "true" ]]; then
        echo "-drive file=$RUN_IMAGE,format=qcow2,if=none,id=disk0,snapshot=on -device ide-hd,drive=disk0"
      else
        echo "-drive file=$RUN_IMAGE,format=qcow2,if=virtio,snapshot=on"
      fi) \
      -nic user,model=$(if [[ "$RUN_AHCI" == "true" ]]; then echo e1000; else echo virtio-net-pci; fi) \
      -device virtio-serial-pci \
      -chardev spicevmc,id=vdagent,debug=0,name=vdagent \
      -device virtserialport,chardev=vdagent,name=com.redhat.spice.0
  fi

  # --- build/copy commands ---
  IMAGE_NAME=""
  GENERALIZE=""
  TO_DISK=""
  TO_REMOTE_DISK=""
  YES=false
  OUT_LINK="./result"

  while [[ ''${#} -gt 0 ]]; do
    case "$1" in
      --image)           IMAGE_NAME="$2"; shift 2 ;;
      --generalize)      GENERALIZE="$2"; shift 2 ;;
      --to-disk)         TO_DISK="$2"; shift 2 ;;
      --to-remote-disk)  TO_REMOTE_DISK="$2"; shift 2 ;;
      -y|--yes)          YES=true; shift ;;
      --out-link)        OUT_LINK="$2"; shift 2 ;;
      --help|-h)         usage ;;
      *)                 echo "Unknown option: $1"; usage ;;
    esac
  done

  [[ -z "$IMAGE_NAME" ]] && { echo "Error: --image is required"; usage; }

  if [[ "$COMMAND" == "copy" && -z "$TO_DISK" && -z "$TO_REMOTE_DISK" ]]; then
    echo "Error: copy requires --to-disk or --to-remote-disk"
    usage
  fi

  # parse --generalize key=val,key=val into nix attrs
  GENERALIZE_EXPR=""
  if [[ -n "$GENERALIZE" ]]; then
    NIX_ARGS=""
    IFS=',' read -ra PAIRS <<< "$GENERALIZE"
    for pair in "''${PAIRS[@]}"; do
      key="''${pair%%=*}"
      val="''${pair#*=}"
      if [[ "$val" == "true" || "$val" == "false" ]]; then
        NIX_ARGS+="$key = $val; "
      else
        NIX_ARGS+="$key = \"$val\"; "
      fi
    done
    # CLI delay-oobe-run → nix delayOobeRun
    NIX_ARGS="''${NIX_ARGS//delay-oobe-run/delayOobeRun}"
    GENERALIZE_EXPR=".generalize { $NIX_ARGS }"
  fi

  FLAKE_DIR="${self}"

  echo "=== vmix $COMMAND ==="
  echo "Image:      $IMAGE_NAME"
  [[ -n "$GENERALIZE" ]] && echo "Generalize: $GENERALIZE"
  [[ -n "$TO_DISK" ]] && echo "To disk:    $TO_DISK"
  [[ -n "$TO_REMOTE_DISK" ]] && echo "To remote:  $TO_REMOTE_DISK"
  echo ""

  # pass DISPLAY+XAUTHORITY to nix builds via temp file (daemon sanitizes env vars)
  # copy Xauthority world-readable so nix build users (nixbld*) can authenticate to X11
  VMIX_DISPLAY_FILE="/tmp/.vmix-display-$$"
  VMIX_XAUTH_FILE="/tmp/.vmix-xauth-$$"
  if [[ -n "''${DISPLAY:-}" ]]; then
    cp "''${XAUTHORITY:-''${HOME}/.Xauthority}" "$VMIX_XAUTH_FILE"
    chmod 644 "$VMIX_XAUTH_FILE"
    printf '%s\n' "''${DISPLAY}" "$VMIX_XAUTH_FILE" > "$VMIX_DISPLAY_FILE"
    chmod 644 "$VMIX_DISPLAY_FILE"
  fi
  trap 'rm -f "$VMIX_DISPLAY_FILE" "$VMIX_XAUTH_FILE"' EXIT

  echo "Building image..."
  ${pkgs.nix}/bin/nix build --out-link "$OUT_LINK" --print-build-logs --impure --option sandbox relaxed --expr "
    let
      vmixLib = (builtins.getFlake \"$FLAKE_DIR\").lib.${system};
      image = vmixLib.$IMAGE_NAME;
    in image$GENERALIZE_EXPR
  "

  IMAGE_FILE=$(readlink -f "$OUT_LINK")
  echo "Built: $IMAGE_FILE"

  # --- write to local disk ---
  if [[ -n "$TO_DISK" ]]; then
    echo ""

    [[ $EUID -ne 0 ]] && { echo "Error: --to-disk requires root"; exit 1; }
    [[ ! -b "$TO_DISK" ]] && { echo "Error: not a block device: $TO_DISK"; exit 1; }

    if ${pkgs.util-linux}/bin/mount | grep -q "^''${TO_DISK}"; then
      echo "Error: $TO_DISK has mounted partitions — unmount first"
      exit 1
    fi

    DISK_SIZE=$(${pkgs.util-linux}/bin/blockdev --getsize64 "$TO_DISK")
    DISK_SIZE_GB=$(( DISK_SIZE / 1024 / 1024 / 1024 ))

    echo "=== writing to disk ==="
    echo "Target: $TO_DISK ($DISK_SIZE_GB GB)"
    echo ""
    echo "WARNING: This will DESTROY all data on $TO_DISK"
    if [[ "$YES" != "true" ]]; then
      read -rp "Continue? [y/N] " confirm
      [[ "$confirm" != "y" && "$confirm" != "Y" ]] && { echo "Aborted."; exit 0; }
    fi

    echo ""
    echo "[1/5] Wiping partition table..."
    ${pkgs.gptfdisk}/bin/sgdisk --zap-all "$TO_DISK"

    echo "[2/5] Writing image to disk..."
    ${pkgs.qemu}/bin/qemu-img convert -p -S 4k -f qcow2 -O raw "$IMAGE_FILE" "$TO_DISK"

    echo "[3/5] Fixing GPT backup header..."
    ${pkgs.gptfdisk}/bin/sgdisk -e "$TO_DISK"
    # delete recovery partition if present, then resize Windows partition
    ${pkgs.gptfdisk}/bin/sgdisk -d 4 "$TO_DISK" 2>/dev/null || true

    echo "[4/5] Expanding Windows partition (partition 3)..."
      ${pkgs.parted}/bin/parted -s "$TO_DISK" resizepart 3 100%

    if [[ "$TO_DISK" == *nvme* ]] || [[ "$TO_DISK" == *mmcblk* ]]; then
      WIN_PART="''${TO_DISK}p3"
    else
      WIN_PART="''${TO_DISK}3"
    fi

    echo "[5/5] Expanding NTFS filesystem on $WIN_PART..."
    ${pkgs.ntfs3g}/bin/ntfsresize --force --no-action "$WIN_PART"
    echo "y" | ${pkgs.ntfs3g}/bin/ntfsresize --force "$WIN_PART"

    echo ""
    echo "Done. $TO_DISK is ready to boot."
  fi

  # --- write to remote disk ---
  if [[ -n "$TO_REMOTE_DISK" ]]; then
    echo ""

    REMOTE_HOST="''${TO_REMOTE_DISK%%:*}"
    REMOTE_DISK="''${TO_REMOTE_DISK#*:}"

    [[ -z "$REMOTE_HOST" || -z "$REMOTE_DISK" ]] && { echo "Error: --to-remote-disk format is user@host:/dev/sdX"; exit 1; }

    echo "=== writing to remote disk ==="
    echo "Host:   $REMOTE_HOST"
    echo "Disk:   $REMOTE_DISK"
    echo ""
    echo "WARNING: This will DESTROY all data on $REMOTE_HOST:$REMOTE_DISK"
    if [[ "$YES" != "true" ]]; then
      read -rp "Continue? [y/N] " confirm
      [[ "$confirm" != "y" && "$confirm" != "Y" ]] && { echo "Aborted."; exit 0; }
    fi

    echo ""
    echo "[1/5] Wiping partition table..."
    ssh "$REMOTE_HOST" "sgdisk --zap-all $REMOTE_DISK"

    echo "[2/5] Streaming image to remote disk..."
    NBD_SOCK="/tmp/vmix-nbd-$$.sock"
    ${pkgs.qemu}/bin/qemu-nbd --read-only -f qcow2 -k "$NBD_SOCK" "$IMAGE_FILE" &
    NBD_PID=$!
    trap "kill $NBD_PID 2>/dev/null; rm -f $NBD_SOCK" EXIT
    while [ ! -S "$NBD_SOCK" ]; do sleep 0.1; done
    DISK_SIZE=$(${pkgs.libnbd}/bin/nbdinfo --size "nbd+unix:///?socket=$NBD_SOCK")
    ${pkgs.libnbd}/bin/nbdcopy --request-size=4194304 "nbd+unix:///?socket=$NBD_SOCK" - \
      | ${pkgs.pv}/bin/pv -s "$DISK_SIZE" \
      | ${pkgs.lib.getBin pkgs.lz4}/bin/lz4 -1 - \
      | ssh "$REMOTE_HOST" "nix-shell -p lz4 --run 'lz4 -d - - | dd of=$REMOTE_DISK bs=4M iflag=fullblock oflag=direct conv=sparse'"
    kill $NBD_PID 2>/dev/null || true
    rm -f "$NBD_SOCK"

    if [[ "$REMOTE_DISK" == *nvme* ]] || [[ "$REMOTE_DISK" == *mmcblk* ]]; then
      REMOTE_WIN_PART="''${REMOTE_DISK}p3"
    else
      REMOTE_WIN_PART="''${REMOTE_DISK}3"
    fi

    echo "[3/5] Fixing GPT backup header..."
    ssh "$REMOTE_HOST" "nix-shell -p gptfdisk --run 'sgdisk -e $REMOTE_DISK && sgdisk -d 4 $REMOTE_DISK 2>/dev/null || true'"

    echo "[4/5] Expanding Windows partition (partition 3)..."
    ssh "$REMOTE_HOST" "nix-shell -p parted --run 'parted -s $REMOTE_DISK resizepart 3 100%'"

    echo "[5/5] Expanding NTFS filesystem on $REMOTE_WIN_PART..."
    ssh "$REMOTE_HOST" "nix-shell -p ntfs3g --run 'echo y | ntfsresize --force $REMOTE_WIN_PART'"

    echo ""
    echo "Done. $REMOTE_HOST:$REMOTE_DISK is ready to boot."
  fi
''