#!/bin/sh

# Cloud Hypervisor MicroVM Management CLI
# Usage: ./ch-vm-cli.sh [list|start|stop|status] [vm_name]

SOCKET_DIR="/var/lib/kvm_agent/sockets"
SYSTEMD_SERVICE_SUFFIX=".service"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Print colored output
print_status() {
    local color="$1"
    local message="$2"
    printf "${color}%s${NC}\n" "$message"
}

# Check if service is active
is_service_active() {
    local service_name="$1"
    if systemctl is-active --quiet "$service_name" 2>/dev/null; then
        return 0 # Service is active
    else
        return 1 # Service is not active or doesn't exist
    fi
}

# Check if VM is running via API
is_vm_running() {
    local socket_path="$1"
    local status
    
    # Check if socket file exists
    if [ ! -S "$socket_path" ]; then
        return 1
    fi
    
    status=$(ch-remote --api-socket "${socket_path}" info 2>/dev/null | jq -r ".state")
    local lower_status=$(echo "$status" | tr '[:upper:]' '[:lower:]')
    
    case "$lower_status" in
        "running"|"active"|"created")
            return 0 # VM is active/created/running
            ;;
        *)
            return 1 # VM is not active
            ;;
    esac
}

# Check if socket is responsive
is_socket_running() {
    local socket_path="$1"
    
    # Check if socket file exists
    if [ ! -S "$socket_path" ]; then
        return 1
    fi
    
    ch-remote --api-socket "${socket_path}" ping >/dev/null 2>&1
    if [ $? -eq 0 ]; then
        return 0 # Socket is responsive
    else
        return 1 # Socket is not responsive
    fi
}

# Get VM status with simplified active check
get_vm_status() {
    local vm_name="$1"
    local socket_name=$(echo "$vm_name" | sed 's/^mvm_//')
    local socket_path="${SOCKET_DIR}/${socket_name}.sock"
    local service_name="${vm_name}${SYSTEMD_SERVICE_SUFFIX}"
    local vm_state=""
    local service_status
    local vm_cpu=""
    local vm_ram=""
    local vm_disk=""

    service_status=$(systemctl is-active "$service_name" 2>/dev/null || echo "inactive")

    # Determine VM state
    if [ "$service_status" != "active" ]; then
        vm_state="inactive"
    else
        if is_socket_running "$socket_path"; then
            local ch_info
            ch_info=$(ch-remote --api-socket "$socket_path" info 2>/dev/null)
            
            local ch_state
            ch_state=$(echo "$ch_info" | jq -r ".state" 2>/dev/null || echo "")
            [ "$ch_state" = "Running" ] && vm_state="active"

            vm_cpu=$(echo "$ch_info" | jq -r ".config.cpus.boot_vcpus" 2>/dev/null)
            # Convert memory from bytes to MB
            local mem_bytes
            mem_bytes=$(echo "$ch_info" | jq -r ".memory_actual_size" 2>/dev/null)
            [ -n "$mem_bytes" ] && vm_ram=$((mem_bytes / 1024 / 1024))MB

            # Get disk paths (comma-separated if multiple)
            vm_disk=$(echo "$ch_info" | jq -r '[.config.disks[].path] | join(", ")' 2>/dev/null)
        fi
    fi

    printf "%-20s %-10s %-10s %-10s %-15s %-15s\n" "$vm_name" "$service_status" "$vm_state" "$vm_cpu" "$vm_ram" "$vm_disk"
}


# List all VMs
list_vms() {
    printf "${BLUE}%-20s %-10s %-10s %-10s %-15s %-15s${NC}\n" "VM NAME" "SERVICE" "STATE" "CPUS" "RAM" "DISKS"
    printf "%-20s %-10s %-10s %-10s %-15s %-15s\n" "-------" "-------" "-----" "----" "----" "-----"
 
    # Find all systemd services that begin with "mvm_"
    systemctl list-units --type=service --all "mvm_*${SYSTEMD_SERVICE_SUFFIX}" --no-pager --no-legend 2>/dev/null | while read -r service_line; do
        service_full=$(echo "$service_line" | awk '{print $1}')
        vm_name=$(echo "$service_full" | sed "s/${SYSTEMD_SERVICE_SUFFIX}$//")
        
        # Only process if it starts with mvm_
        case "$vm_name" in
            mvm_*)
                get_vm_status "$vm_name"
                ;;
        esac
    done
}

# Start a VM
start_vm() {
    local vm_name="$1"
    local service_name="${vm_name}${SYSTEMD_SERVICE_SUFFIX}"
    local socket_name=$(echo "$vm_name" | sed 's/^mvm_//')
    local socket_path="${SOCKET_DIR}/${socket_name}.sock"
    
    if [ -z "$vm_name" ]; then
        print_status "$RED" "Error: VM name is required"
        return 1
    fi
    
    # Check if service exists
    if ! systemctl list-unit-files "$service_name" >/dev/null 2>&1; then
        print_status "$RED" "Error: Service $service_name does not exist"
        return 1
    fi
    
    # Check if service is active
    if is_service_active "$service_name"; then
        print_status "$YELLOW" "Service $service_name is already active"
        
        # Check VM state via ch-remote
        if is_socket_running "$socket_path"; then
            local vm_state=$(ch-remote --api-socket "${socket_path}" info 2>/dev/null | jq -r ".state" 2>/dev/null || echo "unknown")
            
            case "$vm_state" in
                "Running")
                    print_status "$GREEN" "VM $vm_name is already running"
                    return 0
                    ;;
                "Created")
                    print_status "$YELLOW" "VM is in Created state, booting..."
                    if ch-remote --api-socket "${socket_path}" boot 2>/dev/null; then
                        print_status "$GREEN" "VM $vm_name boot command sent successfully"
                    else
                        print_status "$RED" "Failed to boot VM $vm_name"
                        return 1
                    fi
                    ;;
                *)
                    print_status "$YELLOW" "VM $vm_name is in state: $vm_state"
                    ;;
            esac
        else
            print_status "$YELLOW" "Service active but socket not responsive"
        fi
    else
        print_status "$YELLOW" "Starting VM service: $vm_name"
        
        # Start systemd service
        if systemctl start "$service_name"; then
            print_status "$GREEN" "Service $service_name started successfully"
            
            # Wait for socket to become available
            local timeout=10
            local count=0
            
            while [ $count -lt $timeout ]; do
                if is_socket_running "$socket_path"; then
                    local vm_state=$(ch-remote --api-socket "${socket_path}" info 2>/dev/null | jq -r ".state" 2>/dev/null || echo "unknown")
                    
                    if [ "$vm_state" = "Created" ]; then
                        print_status "$YELLOW" "VM created, booting..."
                        if ch-remote --api-socket "${socket_path}" boot 2>/dev/null; then
                            print_status "$GREEN" "VM $vm_name boot command sent successfully"
                        else
                            print_status "$RED" "Failed to boot VM $vm_name"
                            return 1
                        fi
                        return 0
                    elif [ "$vm_state" = "Running" ]; then
                        print_status "$GREEN" "VM $vm_name is now running"
                        return 0
                    fi
                fi
                sleep 1
                count=$((count + 1))
            done
            
            print_status "$YELLOW" "VM service started but may still be initializing"
        else
            print_status "$RED" "Failed to start service $service_name"
            return 1
        fi
    fi
}

# Stop a VM (simplified)
stop_vm() {
    local vm_name="$1"
    local service_name="${vm_name}${SYSTEMD_SERVICE_SUFFIX}"

    if [ -z "$vm_name" ]; then
        print_status "$RED" "Error: VM name is required"
        return 1
    fi

    # Check if service exists
    if ! systemctl list-unit-files "$service_name" >/dev/null 2>&1; then
        print_status "$RED" "Error: Service $service_name does not exist"
        return 1
    fi

    local service_status
    service_status=$(systemctl is-active "$service_name" 2>/dev/null || echo "inactive")

    case "$service_status" in
        active)
            print_status "$YELLOW" "Stopping VM service: $service_name"
            if systemctl stop "$service_name"; then
                print_status "$GREEN" "Service $service_name stopped successfully"
            else
                print_status "$RED" "Failed to stop service $service_name"
                return 1
            fi
            ;;
        inactive)
            print_status "$YELLOW" "Service $service_name is already stopped"
            ;;
        failed)
            print_status "$YELLOW" "Service $service_name is in failed state, resetting..."
            if systemctl reset-failed "$service_name"; then
                print_status "$GREEN" "Service $service_name reset successfully"
            else
                print_status "$RED" "Failed to reset service $service_name"
                return 1
            fi
            ;;
        *)
            print_status "$YELLOW" "Service $service_name is in unknown state: $service_status"
            ;;
    esac
}


# Show status of a specific VM
status_vm() {
    local vm_name="$1"
    
    if [ -z "$vm_name" ]; then
        print_status "$RED" "Error: VM name is required"
        return 1
    fi
    
    # Updated header for 6 columns
    printf "${BLUE}%-20s %-10s %-10s %-10s %-15s %-15s${NC}\n" "VM NAME" "SERVICE" "STATE" "CPUS" "RAM" "DISKS"
    printf "%-20s %-10s %-10s %-10s %-15s %-15s\n" "-------" "-------" "-----" "----" "----" "-----"

    get_vm_status "$vm_name"
}


# Show usage information
show_usage() {
    cat << EOF
Cloud Hypervisor MicroVM Management CLI

Usage: $0 [COMMAND] [VM_NAME]

Commands:
    list                List all VMs and their status
    start <vm_name>     Start a specific VM
    stop <vm_name>      Stop a specific VM  
    status <vm_name>    Show status of a specific VM
    help                Show this help message

Examples:
    $0 list
    $0 start my-vm
    $0 stop my-vm
    $0 status my-vm

Configuration:
    Socket directory: $SOCKET_DIR
    Service pattern: mvm_<vm_name>.service
    Socket pattern: <vm_name>.sock (without mvm_ prefix)

Status Legend:
    SERVICE: systemd service status (active/inactive/failed/etc.)
    VM STATE: VM state from ch-remote info command

EOF
}

# Main function
main() {
    local command="$1"
    local vm_name="$2"
    
    # Check if running as root for systemctl operations
    if [ "$(id -u)" -ne 0 ] && [ "$command" != "list" ] && [ "$command" != "status" ] && [ "$command" != "help" ]; then
        print_status "$RED" "Warning: You may need root privileges to start/stop services"
    fi
    
    # Check if required tools are available
    if ! command -v ch-remote >/dev/null 2>&1; then
        print_status "$RED" "Error: ch-remote command not found. Please install Cloud Hypervisor."
        exit 1
    fi
    
    if ! command -v jq >/dev/null 2>&1; then
        print_status "$RED" "Error: jq command not found. Please install jq."
        exit 1
    fi
    
    case "$command" in
        "list"|"ls")
            list_vms
            ;;
        "start")
            start_vm "$vm_name"
            ;;
        "stop")
            stop_vm "$vm_name"
            ;;
        "status"|"stat")
            status_vm "$vm_name"
            ;;
        "help"|"-h"|"--help")
            show_usage
            ;;
        "")
            show_usage
            exit 1
            ;;
        *)
            print_status "$RED" "Error: Unknown command '$command'"
            show_usage
            exit 1
            ;;
    esac
}

# Execute main function with all arguments
main "$@"