#!/usr/bin/env python3

import os
import sys
import subprocess
import tempfile
import threading

def create_test_fifo():
    """Create a test FIFO and write a passphrase to it."""
    fifo_path = os.path.join(tempfile.gettempdir(), "test_luks_fifo")
    try:
        os.mkfifo(fifo_path)
    except FileExistsError:
        pass

    # Write a test passphrase in a separate thread
    def write_passphrase():
        with open(file=fifo_path, mode="w", encoding="utf-8") as f:
            f.write("test\n")

    threading.Thread(target=write_passphrase, daemon=True).start()
    return fifo_path

def create_test_disk(size_mb=100, fs_type="ext4"):
    """
    Create a test disk image, set up a loop device, and create a filesystem.

    Args:
        size_mb (int): Size of the test disk in megabytes (default: 100).
        fs_type (str): Filesystem type to create (default: "ext4").

    Returns:
        tuple: (loop_device, disk_image_path) or (None, None) if failed.
    """
    # Create a temporary disk image file
    disk_image_path = os.path.join(tempfile.gettempdir(), "test_disk.img")

    try:
        # 1. Create a disk image filled with zeros
        print(f"Creating {size_mb}MB disk image at {disk_image_path}...")
        subprocess.run(
            ["dd", "if=/dev/zero", f"of={disk_image_path}", "bs=1M", f"count={size_mb}"],
            check=True,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )

        # 2. Set up a loop device
        print("Setting up loop device...")
        subprocess.run(
            ["losetup", "-fP", disk_image_path],
            check=True
        )

        # Find the loop device associated with the disk image
        result = subprocess.run(
            ["losetup", "-j", disk_image_path],
            capture_output=True,
            text=True,
            check=True
        )

        # Parse the output to extract the loop device (e.g., /dev/loop0)
        loop_device = None
        for line in result.stdout.splitlines():
            if disk_image_path in line:
                loop_device = line.split(":")[0]
                break

        if not loop_device:
            print("Failed to find loop device")
            os.unlink(disk_image_path)
            return None, None

        print(f"Loop device created: {loop_device}")

        # 3. Create a filesystem and capture output
        print(f"Creating {fs_type} filesystem...")
        mkfs_result = subprocess.run(
            [f"mkfs.{fs_type}", loop_device],
            capture_output=True,
            text=True,
            check=True
        )

        # Print mkfs output for debugging
        print("mkfs output:")
        print(mkfs_result.stdout)
        if mkfs_result.returncode != 0:
            print(f"Failed to create {fs_type} filesystem: {mkfs_result.stderr}")
            subprocess.run(["losetup", "-d", loop_device], capture_output=True, check=True)
            os.unlink(disk_image_path)
            return None, None

        return loop_device, disk_image_path

    except subprocess.CalledProcessError as e:
        print(f"Error: {e}")
        # Clean up if something went wrong
        if 'disk_image_path' in locals() and os.path.exists(disk_image_path):
            os.unlink(disk_image_path)
        return None, None

def cleanup(loop_device, disk_image_path):
    """Clean up the loop device and disk image."""
    if loop_device:
        print(f"Detaching loop device {loop_device}...")
        subprocess.run(["losetup", "-d", loop_device], check=True)

    if disk_image_path and os.path.exists(disk_image_path):
        print(f"Removing disk image {disk_image_path}...")
        os.unlink(disk_image_path)

def debug_script(loop_device):
    """Debug the encryption_sudo.py script."""
    fifo_path = create_test_fifo()

    if os.path.exists(fifo_path):
        print(f"Created test FIFO: {fifo_path}")
    else:
        print(f"Failed to create test FIFO: {fifo_path}")
        sys.exit(1)

    # Call the script with test arguments
    command = [
        "python3",
        "encryption_sudo.py",
        loop_device,
        fifo_path
    ]

    print(f"Running command: {' '.join(command)}")

    # Run the script
    try:
        result = subprocess.run(
            command,
            capture_output=True,
            text=True,
            check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error: {e}")
        sys.exit(1)

    # Print output for debugging
    print("STDOUT:", result.stdout)
    if result.returncode != 0:
        print("STDERR:", result.stderr)
    print("Return code:", result.returncode)

    # Clean up
    try:
        os.unlink(fifo_path)
    except OSError:
        pass

if __name__ == "__main__":
    if os.geteuid() != 0:
        print("This script must be run as root.")
        sys.exit(1)

    USE_LOOP_DEVICE = False

    if not USE_LOOP_DEVICE:
        device_path = "/dev/sdb1"
        debug_script(device_path)
    else:
        # Create test disk
        loop_dev, disk_image = create_test_disk(size_mb=100, fs_type="ext4")

        if loop_dev and disk_image:
            print("\nSuccess! Test disk created:")
            print(f"  Loop device: {loop_dev}")
            print(f"  Disk image:  {disk_image}")

            debug_script(loop_dev)

            # Uncomment to automatically clean up:
            cleanup(loop_dev, disk_image)
        else:
            print("Failed to create test disk.")
