import subprocess
import logging

logger = logging.getLogger(__name__)

def run_command(cmd, check=True):
    logger.info(f"Running command: {' '.join(cmd)}")
    try:
        result = subprocess.run(cmd, check=check, capture_output=True, text=True)
        return result
    except subprocess.CalledProcessError as e:
        logger.error(f"Command failed: {e.stderr}")
        raise

def get_partition_device(disk_device, partition_number):
    """
    Returns the partition device path.
    Handles NVMe style (p1) vs sd style (1).
    """
    if disk_device[-1].isdigit():
        return f"{disk_device}p{partition_number}"
    return f"{disk_device}{partition_number}"

def auto_partition_disk(disk_device):
    """
    Automatically partitions the disk with a standard layout:
    1. EFI System Partition (1GB)
    2. Swap (4GB) - simpler fixed size for now
    3. Root (Remaining)
    """
    logger.info(f"Starting auto-partitioning on {disk_device}")

    # 1. Zap the disk (destroy all data)
    run_command(["sgdisk", "-Z", disk_device])
    
    # 2. Create new GPT table
    run_command(["sgdisk", "-o", disk_device])

    # 3. Create EFI Partition (Part 1, 1GB, Type EF00)
    # -n <partnum>:<start>:<end>
    run_command(["sgdisk", "-n", "1:0:+1024M", "-t", "1:ef00", "-c", "1:EFI System", disk_device])

    # 4. Create Swap Partition (Part 2, 4GB, Type 8200)
    run_command(["sgdisk", "-n", "2:0:+4096M", "-t", "2:8200", "-c", "2:Swap", disk_device])

    # 5. Create Root Partition (Part 3, Rest, Type 8300)
    run_command(["sgdisk", "-n", "3:0:0", "-t", "3:8300", "-c", "3:Root", disk_device])

    # Inform kernel of changes
    run_command(["partprobe", disk_device])
    
    # Wait a bit for nodes to appear? Usually partprobe handles it but sometimes there's a race.
    import time
    time.sleep(1)

    # 6. Format Partitions
    efi_part = get_partition_device(disk_device, 1)
    swap_part = get_partition_device(disk_device, 2)
    root_part = get_partition_device(disk_device, 3)

    logger.info("Formatting EFI partition...")
    run_command(["mkfs.vfat", "-F32", efi_part])

    logger.info("Formatting Swap partition...")
    run_command(["mkswap", swap_part])

    logger.info("Formatting Root partition...")
    run_command(["mkfs.ext4", "-F", root_part])

    logger.info("Partitioning and formatting complete.")
    
    return {
        "efi": efi_part,
        "swap": swap_part,
        "root": root_part
    }