#!/bin/bash

# Disk management module for bootc image creation
# This module handles disk image creation, partitioning, and filesystem setup

# Common functions are sourced by the main script

# Create and partition disk image
create_disk_image_work() {
    log_info "Creating disk image..."
    
    # Check disk space first (temporarily disabled for debugging)
    # check_disk_space
    
    # Create raw image first to ensure proper sizing
    log_info "Creating raw image with size: ${IMAGE_SIZE_GB}GB"
    dd if=/dev/zero of="$OUTPUT_DIR/debian-bootc.raw" bs=1M count=$((IMAGE_SIZE_GB * 1024))
    
    # Create loop device
    loop_dev=$(sudo losetup --find --show "$OUTPUT_DIR/debian-bootc.raw")
    log_info "Using loop device: $loop_dev"
    
    # Export the loop device variable so other modules can use it
    export loop_dev
    
    # Partition the disk with proper GPT layout including BIOS Boot Partition
    log_info "Partitioning disk..."
    sudo parted "$loop_dev" mklabel gpt
    # Create 1MB BIOS Boot Partition for GRUB
    sudo parted "$loop_dev" mkpart primary 1MiB 2MiB
    sudo parted "$loop_dev" set 1 bios_grub on
    # Create main system partition
    sudo parted "$loop_dev" mkpart primary ext4 2MiB 100%
    sudo parted "$loop_dev" set 2 boot on
    
    # Verify partitions were created correctly
    log_info "Verifying partition layout..."
    sudo parted "$loop_dev" print
    
    # Create filesystem on the main partition (partition 2)
    log_info "Creating filesystem..."
    sudo mkfs.ext4 "${loop_dev}p2"
    
    # Mount the main partition (partition 2)
    log_info "Mounting partition..."
    mkdir -p "$WORK_DIR/mnt"
    sudo mount "${loop_dev}p2" "$WORK_DIR/mnt"
    
    log_info "Device size: $(sudo blockdev --getsize64 $loop_dev) bytes"
}

# Create essential directories
create_essential_directories() {
    log_info "Creating essential directories..."
    
    # Create essential directories individually
    sudo mkdir -p "$WORK_DIR/mnt/usr"
    sudo mkdir -p "$WORK_DIR/mnt/etc"
    sudo mkdir -p "$WORK_DIR/mnt/var"
    sudo mkdir -p "$WORK_DIR/mnt/home"
    sudo mkdir -p "$WORK_DIR/mnt/root"
    sudo mkdir -p "$WORK_DIR/mnt/boot"
    sudo mkdir -p "$WORK_DIR/mnt/proc"
    sudo mkdir -p "$WORK_DIR/mnt/sys"
    sudo mkdir -p "$WORK_DIR/mnt/dev"
    sudo mkdir -p "$WORK_DIR/mnt/tmp"
    sudo mkdir -p "$WORK_DIR/mnt/run"
    sudo mkdir -p "$WORK_DIR/mnt/media"
    sudo mkdir -p "$WORK_DIR/mnt/mnt"
    
    log_info "Essential directories created"
}

# Unmount and cleanup disk
cleanup_disk() {
    log_info "Cleaning up disk operations..."
    
    if [ -n "$WORK_DIR" ] && [ -d "$WORK_DIR/mnt" ]; then
        sudo umount "$WORK_DIR/mnt" 2>/dev/null || true
    fi
    
    if [ -n "$loop_dev" ]; then
        sudo losetup -d "$loop_dev" 2>/dev/null || true
    fi
}

# Get disk information
get_disk_info() {
    if [ -n "$loop_dev" ]; then
        log_info "Disk information:"
        sudo parted "$loop_dev" print
        echo ""
        sudo fdisk -l "$loop_dev"
    else
        log_warn "No loop device available for disk info"
    fi
}

# Verify disk integrity
verify_disk_integrity() {
    log_info "Verifying disk integrity..."
    
    if [ ! -f "$OUTPUT_DIR/debian-bootc.raw" ]; then
        log_error "Raw disk image not found"
        return 1
    fi
    
    local image_size=$(stat -c%s "$OUTPUT_DIR/debian-bootc.raw")
    local expected_size=$((IMAGE_SIZE_GB * 1024 * 1024 * 1024))
    
    if [ "$image_size" -ne "$expected_size" ]; then
        log_error "Disk image size mismatch. Expected: $expected_size, Got: $image_size"
        return 1
    fi
    
    log_info "Disk integrity verified successfully"
    return 0
}

# Print module usage
print_module_usage() {
    echo "Usage: source modules/disk.sh"
    echo "This module handles disk image creation and management."
    echo ""
    echo "Available functions:"
    echo "  create_disk_image - Create and partition disk image"
    echo "  create_essential_directories - Create required directory structure"
    echo "  cleanup_disk - Unmount and cleanup disk operations"
    echo "  get_disk_info - Display disk partition information"
    echo "  verify_disk_integrity - Verify disk image integrity"
    echo ""
    echo "Required variables:"
    echo "  OUTPUT_DIR - Output directory for images"
    echo "  IMAGE_SIZE_GB - Image size in gigabytes"
    echo "  WORK_DIR - Working directory for operations"
}