#!/bin/bash

# Proxmox Image Converter
# This script converts various virtual machine image formats to Proxmox-compatible format
# Supported formats: VMDK, VHD, VHDX, QED, QCOW2, RAW

set -euo pipefail

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

# Required tools
REQUIRED_TOOLS=("qemu-img" "md5sum" "awk" "grep" "basename" "dirname")

# Log file
LOG_FILE="/var/log/proxmox_converter.log"
TEMP_DIR="/tmp/proxmox_convert"

# Function to log messages
log_message() {
    local level=$1
    local message=$2
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo -e "${timestamp} [${level}] ${message}" >> "${LOG_FILE}"
    
    case ${level} in
        "ERROR")
            echo -e "${RED}ERROR: ${message}${NC}" >&2
            ;;
        "WARNING")
            echo -e "${YELLOW}WARNING: ${message}${NC}"
            ;;
        "INFO")
            echo -e "${GREEN}INFO: ${message}${NC}"
            ;;
    esac
}

# Function to check required tools
check_requirements() {
    local missing_tools=()
    
    for tool in "${REQUIRED_TOOLS[@]}"; do
        if ! command -v "$tool" >/dev/null 2>&1; then
            missing_tools+=("$tool")
        fi
    done
    
    if [ ${#missing_tools[@]} -ne 0 ]; then
        log_message "ERROR" "Missing required tools: ${missing_tools[*]}"
        log_message "ERROR" "Please install them using: apt-get install qemu-utils coreutils grep"
        exit 1
    fi
}

# Function to verify disk image integrity
verify_image() {
    local image_path=$1
    local checksum_before=$2
    
    if [ ! -f "$image_path" ]; then
        log_message "ERROR" "Image file not found: $image_path"
        return 1
    }
    
    local checksum_after=$(md5sum "$image_path" | awk '{print $1}')
    
    if [ "$checksum_before" != "$checksum_after" ]; then
        log_message "ERROR" "Image verification failed! Checksums don't match"
        return 1
    fi
    
    log_message "INFO" "Image verification successful"
    return 0
}

# Function to detect image format
detect_format() {
    local image_path=$1
    
    if [ ! -f "$image_path" ]; then
        log_message "ERROR" "Image file not found: $image_path"
        return 1
    }
    
    local format=$(qemu-img info "$image_path" | grep "file format" | awk '{print $3}')
    
    if [ -z "$format" ]; then
        log_message "ERROR" "Unable to detect image format"
        return 1
    fi
    
    echo "$format"
    return 0
}

# Function to convert image to QCOW2
convert_to_qcow2() {
    local input_path=$1
    local output_path=$2
    local format=$3
    
    # Calculate checksum before conversion
    local checksum_before=$(md5sum "$input_path" | awk '{print $1}')
    
    # Create temp directory if it doesn't exist
    mkdir -p "$TEMP_DIR"
    
    # Convert image
    log_message "INFO" "Converting $input_path to QCOW2 format..."
    if ! qemu-img convert -f "$format" -O qcow2 "$input_path" "${TEMP_DIR}/temp.qcow2"; then
        log_message "ERROR" "Conversion failed"
        rm -f "${TEMP_DIR}/temp.qcow2"
        return 1
    fi
    
    # Verify source image hasn't been modified
    if ! verify_image "$input_path" "$checksum_before"; then
        log_message "ERROR" "Source image was modified during conversion"
        rm -f "${TEMP_DIR}/temp.qcow2"
        return 1
    fi
    
    # Move converted image to final destination
    mv "${TEMP_DIR}/temp.qcow2" "$output_path"
    
    log_message "INFO" "Conversion completed successfully"
    return 0
}

# Function to import image to Proxmox
import_to_proxmox() {
    local image_path=$1
    local vm_id=$2
    local storage=$3
    
    # Verify VM ID is valid
    if ! [[ "$vm_id" =~ ^[0-9]+$ ]]; then
        log_message "ERROR" "Invalid VM ID: $vm_id"
        return 1
    fi
    
    # Check if VM ID already exists
    if qm status "$vm_id" >/dev/null 2>&1; then
        log_message "ERROR" "VM ID $vm_id already exists"
        return 1
    }
    
    # Check if storage exists
    if ! pvesm status | grep -q "^$storage"; then
        log_message "ERROR" "Storage $storage not found"
        return 1
    }
    
    # Import disk
    log_message "INFO" "Importing disk to Proxmox..."
    if ! qm importdisk "$vm_id" "$image_path" "$storage"; then
        log_message "ERROR" "Failed to import disk"
        return 1
    }
    
    log_message "INFO" "Disk imported successfully"
    return 0
}

# Main function
main() {
    local input_path=$1
    local vm_id=$2
    local storage=$3
    
    # Check requirements
    check_requirements
    
    # Create log directory if it doesn't exist
    mkdir -p "$(dirname "$LOG_FILE")"
    
    # Detect input format
    local format
    format=$(detect_format "$input_path")
    if [ $? -ne 0 ]; then
        exit 1
    fi
    
    # Generate output path
    local filename=$(basename "$input_path")
    local output_path="${TEMP_DIR}/${filename%.*}.qcow2"
    
    # Convert image if not already in QCOW2 format
    if [ "$format" != "qcow2" ]; then
        if ! convert_to_qcow2 "$input_path" "$output_path" "$format"; then
            exit 1
        fi
    else
        cp "$input_path" "$output_path"
    fi
    
    # Import to Proxmox
    if ! import_to_proxmox "$output_path" "$vm_id" "$storage"; then
        rm -f "$output_path"
        exit 1
    fi
    
    # Cleanup
    rm -f "$output_path"
    log_message "INFO" "Image conversion and import completed successfully"
}

# Script usage
usage() {
    echo "Usage: $0 <input_image_path> <vm_id> <storage>"
    echo "Example: $0 /path/to/windows.vmdk 100 local-lvm"
    exit 1
}

# Check arguments
if [ $# -ne 3 ]; then
    usage
fi

# Run main function
main "$1" "$2" "$3"