#!/usr/bin/python3
import argparse
import logging
from ddlib import calculate_md5, checkisomd5, copy_file_to_device, sgdisk, gdisk, get_partition_info, zero_and_verify, is_gpt, print_it

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def parse_arguments():
    parser = argparse.ArgumentParser(description='Process some integers.')
    parser.add_argument('source_file', type=str, help='Source file')
    parser.add_argument('destination_device', type=str, help='Destination device')
    parser.add_argument('block_size', nargs='?', type=str, default='8M', help='Block size (optional, default: 8M)')
    return parser.parse_args()

def main():
    args = parse_arguments()
    
    if not zero_and_verify(args.destination_device):
        logger.error("Zero check failed")
        return
    
    if not checkisomd5(args.source_file):
        logger.error("ISO MD5 checksum verification failed")
        return
    
    params = {
        'source_file': args.source_file,
        'destination_device': args.destination_device,
        'block_size': args.block_size
    }
    
    if not copy_file_to_device(params):
        logger.error("Copy file to device failed")
        return
    
    if is_gpt(args.destination_device):
        logger.info("Aligning GPT")
        sgdisk(args.destination_device, "-e", custom_print=print_it)
        
        logger.info("Creating hybrid partition (GPT/MBR)")
        info = get_partition_info(args.destination_device)
        num_parts = ", ".join(str(i) for i in range(1, len(info) + 1))
        sgdisk(args.destination_device, {
            "-h": num_parts,
            "-U": "random"
        }, custom_print=print_it)
    
    logger.info("Verification")
    print(gdisk(args.destination_device, ["-l", ], custom_print=print_it).stdout)

if __name__ == "__main__":
    main()
