#!/usr/bin/env bash
# proxmox-screenshot -- Take a console screenshot of a Proxmox VM
#
# Usage: proxmox-screenshot VMID [OUTPUT.png]
#
# Requires: sshpass, python3-pil (for PNG conversion)
# Config:   proxmox.yaml (host), env file (PROXMOX_ROOT_PASSWORD)
#
# Outputs the path to the PNG file on stdout.

set -euo pipefail

# ---------------------------------------------------------------------------
# Usage
# ---------------------------------------------------------------------------

usage() {
    cat <<'EOF'
Usage: proxmox-screenshot VMID [OUTPUT.png]

Take a console screenshot of a Proxmox VM via QEMU monitor screendump.

Arguments:
  VMID          VM ID (must be in allowed range: 400-499, 800-939)
  OUTPUT.png    Output file path (default: /tmp/vm-VMID-screen.png)

Environment:
  PROXMOX_ROOT_PASSWORD   SSH root password (from env file)

Config:
  proxmox.yaml            Proxmox host (auto-discovered)

Examples:
  proxmox-screenshot 900
  proxmox-screenshot 900 /tmp/my-screenshot.png
EOF
    exit 0
}

if [[ "${1:-}" == "--help" || "${1:-}" == "-h" ]]; then
    usage
fi

# ---------------------------------------------------------------------------
# Validate arguments
# ---------------------------------------------------------------------------

if [[ $# -lt 1 ]]; then
    echo "Error: VMID required. Usage: proxmox-screenshot VMID [OUTPUT.png]" >&2
    exit 1
fi

VMID="$1"
OUTPUT="${2:-/tmp/vm-${VMID}-screen.png}"

# VMID safety check: only allowed ranges
if ! [[ "$VMID" =~ ^[0-9]+$ ]]; then
    echo "Error: VMID must be a number" >&2
    exit 1
fi

if ! (( (VMID >= 400 && VMID <= 499) || (VMID >= 800 && VMID <= 939) )); then
    echo "Error: VMID $VMID not in allowed range (400-499, 800-939)" >&2
    exit 1
fi

# ---------------------------------------------------------------------------
# Load environment
# ---------------------------------------------------------------------------

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

load_env() {
    # Already exported (user sourced envs/*/env)
    [[ -n "${PROXMOX_ROOT_PASSWORD:-}" ]] && return

    # Check envs/*/env in repo root
    local repo_root
    repo_root=$(cd "$SCRIPT_DIR/../.." && pwd)
    for env_file in "${repo_root}"/envs/*/env; do
        if [[ -f "$env_file" ]]; then
            # shellcheck disable=SC1091
            source "$env_file"
            [[ -n "${PROXMOX_ROOT_PASSWORD:-}" ]] && return
        fi
    done

    # Walk up from script dir (backward compat)
    local dir="$SCRIPT_DIR"
    while [[ "$dir" != "/" ]]; do
        if [[ -f "${dir}/env" ]]; then
            # shellcheck disable=SC1091
            source "${dir}/env"
            return
        fi
        dir="$(dirname "$dir")"
    done
}
load_env

if [[ -z "${PROXMOX_ROOT_PASSWORD:-}" ]]; then
    echo "Error: PROXMOX_ROOT_PASSWORD not set. Add it to the env file." >&2
    exit 1
fi

# ---------------------------------------------------------------------------
# Load proxmox.yaml for host
# ---------------------------------------------------------------------------

yaml_get() {
    local file="$1" key="$2" default="${3:-}"
    local val
    val=$(awk -F': ' -v k="$key" '$1 == k {print $2; exit}' "$file" 2>/dev/null) || val=""
    val="${val#\"}" ; val="${val%\"}"
    val="${val#\'}" ; val="${val%\'}"
    echo "${val:-$default}"
}

find_proxmox_yaml() {
    local repo_root
    repo_root=$(cd "$SCRIPT_DIR/../.." && pwd)

    # Check envs/*/proxmox.yaml in repo root
    for f in "${repo_root}"/envs/*/proxmox.yaml; do
        if [[ -f "$f" ]]; then
            echo "$f"
            return
        fi
    done

    # Walk up from script dir (backward compat)
    local search_dirs=(
        "${SCRIPT_DIR}/.."
        "${repo_root}"
        "$(pwd)"
    )
    for dir in "${search_dirs[@]}"; do
        if [[ -f "${dir}/proxmox.yaml" ]]; then
            echo "${dir}/proxmox.yaml"
            return
        fi
    done
    echo ""
}

PROXMOX_YAML=$(find_proxmox_yaml)
if [[ -z "$PROXMOX_YAML" ]]; then
    echo "Error: proxmox.yaml not found" >&2
    exit 1
fi

PVE_HOST=$(yaml_get "$PROXMOX_YAML" "host")
if [[ -z "$PVE_HOST" ]]; then
    echo "Error: proxmox.yaml: 'host' is required" >&2
    exit 1
fi

# ---------------------------------------------------------------------------
# Take screenshot
# ---------------------------------------------------------------------------

SSH_OPTS=(-o StrictHostKeyChecking=no -o ConnectTimeout=5 -o LogLevel=ERROR)
REMOTE_PPM="/tmp/vm-${VMID}-screen.ppm"
LOCAL_PPM="/tmp/vm-${VMID}-screen.ppm"

# Take screenshot via QEMU monitor
if ! sshpass -p "$PROXMOX_ROOT_PASSWORD" ssh "${SSH_OPTS[@]}" \
    "root@${PVE_HOST}" \
    "echo 'screendump ${REMOTE_PPM}' | qm monitor ${VMID}" &>/dev/null; then
    echo "Error: screendump failed (VM $VMID not running or SSH error)" >&2
    exit 1
fi

# Retrieve PPM file
if ! sshpass -p "$PROXMOX_ROOT_PASSWORD" scp "${SSH_OPTS[@]}" \
    "root@${PVE_HOST}:${REMOTE_PPM}" "$LOCAL_PPM" &>/dev/null; then
    echo "Error: failed to retrieve screenshot from Proxmox host" >&2
    exit 1
fi

# Cleanup remote file
sshpass -p "$PROXMOX_ROOT_PASSWORD" ssh "${SSH_OPTS[@]}" \
    "root@${PVE_HOST}" "rm -f ${REMOTE_PPM}" &>/dev/null || true

# Convert PPM -> PNG
if python3 -c "from PIL import Image; Image.open('${LOCAL_PPM}').save('${OUTPUT}')" 2>/dev/null; then
    rm -f "$LOCAL_PPM"
else
    # PIL not available, keep as PPM
    OUTPUT="${OUTPUT%.png}.ppm"
    mv "$LOCAL_PPM" "$OUTPUT" 2>/dev/null || true
fi

echo "$OUTPUT"