import os
import subprocess
import shutil
import tempfile
import argparse
import time

def check_dependencies():
    required_tools = ['erofsfuse', 'fusermount', 'mkfs.ext4']
    missing = [tool for tool in required_tools if shutil.which(tool) is None]
    if missing:
        raise RuntimeError(f"Missing tools: {', '.join(missing)}. Install erofs-utils, fuse, and e2fsprogs.")

def is_mounted(mount_point):
    result = subprocess.run(
        ['findmnt', '-rn', '-T', mount_point],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL
    )
    return result.returncode == 0

def create_ext4_image(output_image, size_mb):
    subprocess.run(['dd', 'if=/dev/zero', f'of={output_image}', 'bs=1M', f'count={size_mb}'], check=True)
    subprocess.run(['mkfs.ext4', '-F', output_image], check=True)

def safe_copy(src, dst):
    try:
        if os.path.islink(src):
            linkto = os.readlink(src)
            os.symlink(linkto, dst)
        elif os.path.isfile(src):
            shutil.copy2(src, dst, follow_symlinks=False)
        elif os.path.isdir(src):
            shutil.copytree(src, dst, symlinks=True, dirs_exist_ok=True)
    except (shutil.Error, OSError) as e:
        print(f"Warning: {str(e)} - Skipping {src}")

def convert_erofs_to_ext4(erofs_image, output_image):
    check_dependencies()
    
    with tempfile.TemporaryDirectory() as temp_dir:
        erofs_mount = os.path.join(temp_dir, 'erofs_mount')
        extract_dir = os.path.join(temp_dir, 'extracted_data')
        os.makedirs(erofs_mount, exist_ok=True)
        os.makedirs(extract_dir, exist_ok=True)
        
        fuse_proc = subprocess.Popen(['erofsfuse', erofs_image, erofs_mount])
        try:
            for _ in range(10):
                if is_mounted(erofs_mount):
                    break
                time.sleep(1)
            else:
                raise RuntimeError("EROFS mount timed out")
            
            for item in os.listdir(erofs_mount):
                src = os.path.join(erofs_mount, item)
                dst = os.path.join(extract_dir, item)
                safe_copy(src, dst)
                
        finally:
            try:
                subprocess.run(['fusermount', '-u', erofs_mount], check=True)
            except subprocess.CalledProcessError:
                pass
            fuse_proc.terminate()
            fuse_proc.wait()
        
        data_size = sum(os.path.getsize(os.path.join(dirpath, f)) 
                      for dirpath, _, files in os.walk(extract_dir) 
                      for f in files if os.path.isfile(os.path.join(dirpath, f)))
        size_mb = int((data_size / (1024 * 1024) + 10) * 1.2)
        
        create_ext4_image(output_image, size_mb)
        
        ext4_mount = os.path.join(temp_dir, 'ext4_mount')
        os.makedirs(ext4_mount, exist_ok=True)
        subprocess.run(['mount', '-o', 'loop', output_image, ext4_mount], check=True)
        try:
            subprocess.run(['rm', '-rf', f'{ext4_mount}/*'], check=True)
            subprocess.run(['rsync', '-a', '--safe-links', f'{extract_dir}/', ext4_mount], check=True)
        finally:
            subprocess.run(['umount', ext4_mount], check=True)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Convert EROFS image to ext4 using erofsfuse')
    parser.add_argument('input_erofs', help='Input EROFS image file')
    parser.add_argument('output_ext4', help='Output ext4 image file')
    args = parser.parse_args()
    
    try:
        convert_erofs_to_ext4(args.input_erofs, args.output_ext4)
        print("Completed!")
    except Exception as e:
        print(f"Error: {str(e)}")
        exit(1)