#!/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
    # Handle help flags first
    if [[ $# -gt 0 && ( "$1" == "-h" || "$1" == "--help" ) ]]; then
        show_usage
        exit 0
    fi
    
    # If no command given, default to interactive zsh shell
    local run_shell=false
    if [[ $# -eq 0 ]]; then
        run_shell=true
    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
    
    # If this is a newly registered workspace, restart VM to mount it
    if [[ "$newly_added" -eq 1 ]] && is_vm_running; then
        log_info "New workspace registered. Restarting VM to mount it..."
        
        # Path to qvm-stop and qvm-start scripts
        local script_dir="${QVM_LIB_DIR}/../bin"
        local qvm_stop="$script_dir/qvm-stop"
        local qvm_start="$script_dir/qvm-start"
        
        # Stop the VM
        if ! "$qvm_stop"; then
            die "Failed to stop VM"
        fi
        
        # Start the VM with new workspace mount
        if ! "$qvm_start"; then
            die "Failed to start VM"
        fi
        
        log_info "VM restarted with new workspace mounted"
    fi
    
    # Ensure VM is running (if it wasn't running before)
    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 (or shell)
    local remote_cmd="cd '$guest_path'"
    
    if [[ "$run_shell" == "true" ]]; then
        # No command - start interactive zsh shell
        remote_cmd+=" && exec zsh"
    else
        # Append user's command with proper quoting
        remote_cmd+=" && "
        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
    fi
    
    # 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 "$@"