#!/usr/bin/env bash
#
# qvm-run - Execute a command in the VM workspace
#
# This script:
# 1. Ensures VM is running (auto-starts if needed)
# 2. Registers current $PWD as a workspace in workspaces.json
# 3. SSHes into VM and executes command in workspace mount point
# 4. Streams output and preserves exit code
#
# Usage: qvm-run <command> [args...]
#
# Notes:
# - Workspaces are pre-mounted at VM start time (no dynamic 9p hotplug)
# - If workspace not already registered, warns to restart VM
# - Uses workspace hash for mount tag and path
#

set -euo pipefail

# Source common library
QVM_LIB_DIR="${QVM_LIB_DIR:-$(dirname "$(readlink -f "$0")")/../lib}"
# shellcheck source=lib/common.sh
source "$QVM_LIB_DIR/common.sh"

#
# show_usage - Display usage information
#
show_usage() {
    cat <<EOF
Usage: qvm run <command> [args...]

Execute a command in the VM at the current workspace directory.

The current directory (\$PWD) is automatically registered as a workspace
and mounted into the VM. Commands run in the mounted workspace directory.

Examples:
  qvm run cargo build
  qvm run npm install
  qvm run ls -la
  qvm run bash -c "pwd && ls"

Notes:
  - Workspaces are mounted at VM start time
  - If this workspace is new, you'll need to restart the VM
  - Command output streams to your terminal
  - Exit code matches the command's exit code

EOF
}

#
# register_workspace - Add workspace to registry if not already present
# Args: $1 - absolute path to workspace
#       $2 - workspace hash
# Returns: 0 if already registered, 1 if newly added (requires VM restart)
#
register_workspace() {
    local workspace_path="$1"
    local hash="$2"
    local mount_tag="ws_${hash}"
    local guest_path="/workspace/${hash}"
    
    # Create workspaces.json if it doesn't exist
    if [[ ! -f "$QVM_WORKSPACES_FILE" ]]; then
        echo '[]' > "$QVM_WORKSPACES_FILE"
    fi
    
    # Check if workspace already registered
    if jq -e --arg path "$workspace_path" '.[] | select(.host_path == $path)' "$QVM_WORKSPACES_FILE" >/dev/null 2>&1; then
        log_info "Workspace already registered: $workspace_path"
        return 0
    fi
    
    # Add new workspace to registry
    log_info "Registering new workspace: $workspace_path"
    local temp_file
    temp_file=$(mktemp)
    
    jq --arg path "$workspace_path" \
       --arg hash "$hash" \
       --arg tag "$mount_tag" \
       --arg guest "$guest_path" \
       '. += [{
           host_path: $path,
           hash: $hash,
           mount_tag: $tag,
           guest_path: $guest
       }]' "$QVM_WORKSPACES_FILE" > "$temp_file"
    
    mv "$temp_file" "$QVM_WORKSPACES_FILE"
    log_info "Workspace registered as $mount_tag -> $guest_path"
    
    return 1  # Indicate new workspace added
}

#
# ensure_workspace_mounted - Mount workspace in VM if not already mounted
# Args: $1 - SSH port
#       $2 - mount tag (e.g., ws_abc123)
#       $3 - guest path (e.g., /workspace/abc123)
# Returns: 0 on success
#
ensure_workspace_mounted() {
    local ssh_port="$1"
    local mount_tag="$2"
    local guest_path="$3"
    
    # SSH into VM and mount the workspace
    # - mkdir creates the mount point if missing
    # - mount attempts to mount the 9p virtfs
    # - || true ensures we don't fail if already mounted
    sshpass -p root ssh -o StrictHostKeyChecking=no \
           -o UserKnownHostsFile=/dev/null \
           -o LogLevel=ERROR \
           -o PubkeyAuthentication=no \
           -o PasswordAuthentication=yes \
           -p "$ssh_port" \
           root@localhost \
           "mkdir -p '$guest_path' && mount -t 9p -o trans=virtio,version=9p2000.L,msize=104857600 '$mount_tag' '$guest_path' 2>/dev/null || true" >/dev/null 2>&1
    
    return 0
}

#
# is_workspace_mounted - Check if workspace is actually mounted in VM
# Args: $1 - SSH port
#       $2 - guest path
# Returns: 0 if mounted, 1 if not
#
is_workspace_mounted() {
    local ssh_port="$1"
    local guest_path="$2"
    
    # SSH into VM and check if guest path exists and is a directory
    if sshpass -p root ssh -o StrictHostKeyChecking=no \
           -o UserKnownHostsFile=/dev/null \
           -o LogLevel=ERROR \
           -o PubkeyAuthentication=no \
           -o PasswordAuthentication=yes \
           -p "$ssh_port" \
           root@localhost \
           "test -d '$guest_path'" 2>/dev/null; then
        return 0
    else
        return 1
    fi
}

#
# main - Main execution flow
#
main() {
    # Show usage if no arguments
    if [[ $# -eq 0 ]]; then
        show_usage
        exit 1
    fi
    
    # Handle help flags
    if [[ "$1" == "-h" || "$1" == "--help" ]]; then
        show_usage
        exit 0
    fi
    
    # Get current workspace (absolute path)
    local workspace_path
    workspace_path="$(pwd)"
    
    # Generate workspace hash
    local hash
    hash=$(workspace_hash "$workspace_path")
    
    local guest_path="/workspace/${hash}"
    
    log_info "Workspace: $workspace_path"
    log_info "Guest path: $guest_path"
    
    # Register workspace in registry
    local newly_added=0
    if ! register_workspace "$workspace_path" "$hash"; then
        newly_added=1
    fi
    
    # Ensure VM is running
    if ! is_vm_running; then
        log_info "VM not running, starting..."
        
        # Path to qvm-start script
        local qvm_start="${QVM_LIB_DIR}/../bin/qvm-start"
        
        if ! "$qvm_start"; then
            die "Failed to start VM"
        fi
    fi
    
    # Get SSH port
    local ssh_port
    ssh_port=$(get_ssh_port)
    
    # Get mount tag from workspaces.json
    local mount_tag
    mount_tag=$(jq -r --arg path "$workspace_path" '.[] | select(.host_path == $path) | .mount_tag' "$QVM_WORKSPACES_FILE")
    
    # Ensure workspace is mounted (auto-mount if not)
    log_info "Ensuring workspace is mounted..."
    ensure_workspace_mounted "$ssh_port" "$mount_tag" "$guest_path"
    
    # Verify workspace is actually mounted
    if ! is_workspace_mounted "$ssh_port" "$guest_path"; then
        log_error "Failed to mount workspace in VM"
        echo ""
        echo "The workspace could not be mounted automatically."
        echo "This may indicate the VM was started before this workspace was registered."
        echo ""
        echo "Please restart the VM to properly configure the workspace:"
        echo "  qvm stop"
        echo "  qvm start"
        echo ""
        echo "Then try your command again."
        exit 1
    fi
    
    # Build SSH command
    # - Use sshpass for automated password auth (password: root)
    # - Use -t if stdin is a TTY (for interactive commands)
    # - Suppress SSH warnings (ephemeral VM, host keys change)
    # - cd to guest path and execute command
    local ssh_cmd=(
        sshpass -p root
        ssh
        -o StrictHostKeyChecking=no
        -o UserKnownHostsFile=/dev/null
        -o LogLevel=ERROR
        -o PubkeyAuthentication=no
        -o PasswordAuthentication=yes
        -p "$ssh_port"
    )
    
    # Add -t flag if stdin is a TTY
    if [[ -t 0 ]]; then
        ssh_cmd+=(-t)
    fi
    
    # Add connection target
    ssh_cmd+=(root@localhost)
    
    # Build remote command: cd to workspace and execute user's command
    # Quote each argument properly to handle spaces and special chars
    local remote_cmd="cd '$guest_path' && "
    
    # Append user's command with proper quoting
    local first_arg=1
    for arg in "$@"; do
        if [[ $first_arg -eq 1 ]]; then
            remote_cmd+="$arg"
            first_arg=0
        else
            # Quote arguments that contain spaces or special characters
            if [[ "$arg" =~ [[:space:]] ]]; then
                remote_cmd+=" '$arg'"
            else
                remote_cmd+=" $arg"
            fi
        fi
    done
    
    # Add the remote command as final SSH argument
    ssh_cmd+=("$remote_cmd")
    
    # Execute SSH command (replaces current process)
    exec "${ssh_cmd[@]}"
}

# Run main function
main "$@"