#!/usr/bin/env python3
""" 
Module for getting partition information.

Author: Arjen Balfoort
Date: 2026-02-22
"""

import os
import subprocess
from typing import Optional
from utils import DebSoftwareApp


class PartitionInfo():
    """
    Class for getting information about partitions.
    """
    def __init__(self, app: DebSoftwareApp):
        # Set the global logging handler
        self.app: DebSoftwareApp = app
        # Need to set this first to None
        # or else it will error out in get_path_from_uuid
        self.partitions = None
        self.partitions: list[dict] = self.get_partitions()
        self.fstab_mounts: dict = None

    @staticmethod
    def run_command(command: str) -> list[str]:
        """
        Run a command and return its output as a list of lines.
        """
        result = subprocess.run(command,
                                shell=True,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE,
                                text=True,
                                check=True)
        return result.stdout.splitlines()

    def mount_partition(self, device_path: str) -> Optional[str]:
        """
        Temporarily mounts a partition to a sub-directory in /tmp.
        Use --no-user-interaction to prevent the user from being prompted.

        Args:
            device_path (str): The path to the device/partition (e.g., "/dev/sdb1").
            fs_type (str): The filesystem type (e.g., "ext4").

        Returns:
            str: The path to the temporary mount directory, or None if mounting failed.
        """
        try:
            result = subprocess.run(
                ["udisksctl", "mount", "-b", device_path, "--no-user-interaction"],
                capture_output=True,
                text=True,
                check=False
            )
            if result.returncode == 1:
                self.app.logging.debug("Cannot user mount %s: skipping", device_path)
                return None
            mount_point = result.stdout.strip().split()[-1]
            return mount_point
        except subprocess.CalledProcessError as e:
            self.app.logging.error("Failed to mount %s: %s", device_path, e)
            return None

    def unmount_partition(self, device_path: str) -> bool:
        """
        Unmounts the partition.

        Args:
            device_path (str): The partition device path.
            A mount point is also accepted.

        Returns:
            bool: True if unmounted, False otherwise.
        """
        try:
            if not device_path.startswith("/dev"):
                # Get the device for the mount point
                result = subprocess.run(
                    ["findmnt", "-n", "-o", "SOURCE", device_path],
                    capture_output=True,
                    text=True,
                    check=True
                )
                device_path = result.stdout.strip()

            # Unmount the device
            subprocess.run(
                ["udisksctl", "unmount", "-b", device_path],
                check=True,
                capture_output=True,
                text=True
            )
            return True
        except Exception as e:
            self.app.logging.error("Failed to unmount %s: %s", device_path, e)
            return False

    def refresh_partitions(self):
        """ Refreshes the list of partitions. """
        self.partitions = self.get_partitions()

    def get_partitions(self) -> list[dict]:
        """
        Get information about all partitions.

        Returns:
            list[dict]: A list of dictionaries, where each dictionary
            contains information about a partition:
                PATH,UUID,LABEL,MOUNTPOINT,FSTYPE,FSSIZE,FSAVAIL,HOTPLUG
        """
        partitions = []
        temp_mounts = []

        # Temporarily mount partitions to get all data of those partitions
        lsblk_output = self.run_command("lsblk -o PATH,FSTYPE,MOUNTPOINT -P")
        for line in lsblk_output:
            # Remove double quotes
            line = line.replace('"', '').strip()
            # Create a dictionary from the line
            parts = dict(item.split("=") for item in line.split() if "=" in item)
            if not parts.get("MOUNTPOINT") and \
               parts.get("FSTYPE") and \
               'crypt' not in parts.get("FSTYPE"):
                temp_mount = self.mount_partition(parts.get("PATH"))
                if temp_mount is not None:
                    temp_mounts.append(temp_mount)

        # Get the partition data needed
        lsblk_output = self.run_command(
            "lsblk -o PATH,UUID,LABEL,MOUNTPOINT,FSTYPE,FSSIZE,FSAVAIL,RM,HOTPLUG -P -b"
        )
        for line in lsblk_output:
            # Remove double quotes
            line = line.replace('"', '').strip()
            # Create a dictionary from the line
            parts = dict(item.split("=") for item in line.split() if "=" in item)

            if not parts.get("FSTYPE"):
                # Exclude non-mountable partitions or disks
                continue

            # Add the dictionary to the partitions list
            partitions.append(parts)

        # Unmount the temporary mounts
        for temp_mount in temp_mounts:
            for partition in partitions:
                if partition.get("MOUNTPOINT") == temp_mount:
                    partition["MOUNTPOINT"] = ''
                    break
            self.unmount_partition(temp_mount)

        # Return partitions data
        return self._filter_partitions(partitions)

    def _filter_partitions(self, partitions: list[dict]) -> list[dict]:
        """ Filter the list of partitions """
        non_mountables = [
            '/dev/sr', '/dev/loop', '/dev/zd', '/dev/ram',
            '/dev/dm-', '/dev/ubi', '/dev/mtd'
        ]
        allowed_fs_types = ['ext4', 'xfs', 'btrfs', 'vfat', 'ntfs', 'exfat', 'f2fs', 'crypto_LUKS']

        self.fstab_mounts = self.get_fstab_mounts(partitions)
        filtered_partitions = []
        mapped_partitions = {}

        for partition in partitions:
            device_path = partition.get("PATH")
            underlying_device = ''

            if device_path in mapped_partitions:
                # Get additional information about the underlying partition
                partition["RM"] = mapped_partitions[device_path].get("RM")
                partition["HOTPLUG"] = mapped_partitions[device_path].get("HOTPLUG")
                underlying_device = mapped_partitions[device_path].get("PATH")

            # Skip if device_path starts with any of the non_mountables
            if any(device_path.startswith(prefix) for prefix in non_mountables):
                continue

            # Skip if filesystem type is not in allowed_fs_types
            if partition.get("FSTYPE") not in allowed_fs_types:
                continue

            # Skip if the partition is encrypted and mapped
            mapper_path = f"/dev/mapper/luks-{partition.get('UUID')}"
            if partition.get("FSTYPE") == 'crypto_LUKS' and \
               os.path.exists(mapper_path):
                mapped_partitions[mapper_path] = partition
                continue

            # Exclude fstab configured partitions
            fstab_partition_mounts = self.fstab_mounts.get(device_path, [])
            if fstab_partition_mounts:
                continue

            # Get information about the partition
            info = {}
            mount_point = partition.get("MOUNTPOINT")
            is_encrypted = False
            if underlying_device or partition.get("FSTYPE") == 'crypto_LUKS':
                is_encrypted = True
            info[device_path] = {
                "uuid": partition.get("UUID"),
                "label": partition.get("LABEL"),
                "mount_point": mount_point,
                "fs_type": partition.get("FSTYPE"),
                "fs_size": partition.get("FSSIZE"),
                "fs_avail": partition.get("FSAVAIL"),
                "underlying_device": underlying_device,
                "is_encrypted": is_encrypted,
                "is_usb": bool(partition.get("HOTPLUG") == '1'),
                "is_fstab_mounted": bool(fstab_partition_mounts)
            }

            filtered_partitions.append(info)
        return filtered_partitions

    def get_path_from_uuid(self, device_uuid: str) -> Optional[str]:
        """
        Get the partition path of a given UUID.

        Args:
            uuid: UUID of the partition.

        Returns:
            str: Partition path if found, else empty string.
        """
        if self.partitions is not None:
            for partition in self.partitions:
                if device_uuid == partition.get("uuid"):
                    return partition.keys()[0]

        # Fallback method 1: /dev/disk/by-uuid/
        uuid_path = f"/dev/disk/by-uuid/{device_uuid}"
        if os.path.exists(uuid_path):
            return os.path.realpath(uuid_path)

        # Fallback method 2: /proc/mounts (mounted devices only)
        try:
            with open(file="/proc/mounts", mode="r", encoding="utf-8") as f:
                for line in f:
                    parts = line.split()
                    if len(parts) >= 3 and device_uuid in parts[0]:
                        return parts[0].split("[")[0]
        except IOError:
            pass

        return None

    @staticmethod
    def get_underlying_device(device_path: str) -> Optional[str]:
        """
        Get the underlying device path of a mapped/encrypted device.

        Args:
            device_path: Path to the device (e.g., "/dev/sda1" or "/dev/mapper/luks-UUID").

        Returns:
            Optional[str]: Underlying device path or None if not found.
        """
        is_root = os.geteuid() == 0
        # If the device is a mapped/encrypted device, resolve the underlying device
        if "/dev/mapper/" in device_path:
            if is_root:
                try:
                    result = subprocess.run(
                        ["dmsetup", "table", "--target", "crypt", device_path],
                        capture_output=True,
                        text=True,
                        check=True
                    )
                    underlying_device = result.stdout.split()[-1].strip()
                    device_path = underlying_device
                except subprocess.CalledProcessError:
                    return None
            else:
                return None  # Cannot resolve mapped devices without root
        return device_path

    def get_uuid_from_device(self, device_path: str) -> Optional[str]:
        """
        Get the UUID from a device path.
        If the device is a mapped/encrypted device (e.g., /dev/mapper/luks-UUID),
        the UUID of the underlying device (e.g., /dev/sda1) is returned.

        Args:
            device_path: Path to the device (e.g., "/dev/sda1" or "/dev/mapper/luks-UUID").

        Returns:
            Optional[str]: UUID of the device or underlying device, or None if not found.
        """
        # Get the underlying device if it is a mapped device
        device_path = self.get_underlying_device(device_path)

        if device_path is None:
            return None

        # Try to get UUID with blkid (requires root for most devices)
        if os.geteuid() == 0:
            try:
                result = subprocess.run(
                    ["blkid", "-s", "UUID", "-o", "value", device_path],
                    capture_output=True,
                    text=True,
                    check=True
                )
                return result.stdout.strip()
            except subprocess.CalledProcessError:
                pass  # Fall through to non-root method

        # Fallback: Check /dev/disk/by-uuid/ (no root needed)
        try:
            for uuid in os.listdir("/dev/disk/by-uuid"):
                uuid_path = os.path.join("/dev/disk/by-uuid", uuid)
                if os.path.realpath(uuid_path) == device_path:
                    return uuid
        except FileNotFoundError:
            pass

        return None

    def get_fstab_mounts(self, partitions: dict) -> dict:
        """
        Get mount paths from /etc/fstab.

        Args:
            partitions (list): A list of dictionaries,
            where each dictionary contains information about a partition.

        Returns:
            dict: A dictionary where the keys are the UUIDs of the partitions
            and the values are the mount paths.
        """
        fstab_mounts = {}
        for partition in partitions:
            temp_mount = None
            mount_point = partition.get("MOUNTPOINT")
            if not mount_point and \
               partition.get("FSTYPE") and \
               'crypt' not in partition.get("FSTYPE"):
                # Mount the partition (only if not already mounted)
                mount_point = self.mount_partition(partition.get("PATH"))
                if mount_point:
                    temp_mount = mount_point

            if mount_point is None:
                # Skip non-mountable partitions
                continue

            fstab_path = f"{mount_point}/etc/fstab"
            if os.path.exists(fstab_path):
                with open(file=fstab_path, mode="r", encoding="utf-8") as fstab:
                    for line in fstab:
                        if not line.strip().startswith("#"):
                            parts = line.split()
                            if len(parts) >= 2 and '=' in parts[0]:
                                device_uuid = parts[0].split('=')[1]
                                device_path = self.get_path_from_uuid(device_uuid)
                                if device_path is not None:
                                    # The partition can be fstab mounted multiple times
                                    # on a multi-boot system
                                    try:
                                        mounts = fstab_mounts[device_path]
                                        if parts[1] not in mounts:
                                            fstab_mounts[device_path].append(parts[1])
                                    except KeyError:
                                        fstab_mounts[device_path] = [parts[1]]

            if temp_mount is not None:
                # Unmount the partition
                self.unmount_partition(temp_mount)

        return fstab_mounts

    def get_available_size(self, device_path: str) -> Optional[float]:
        """
        Get the available size of a device in MiB.
        If the device is a mapped/encrypted device (e.g., /dev/mapper/luks-UUID),
        the available size of the underlying device is returned.

        Args:
            device_path (str): The path to the device.

        Returns:
            float: The available size of the device in MiB.
        """
        # Get the underlying device if it is a mapped device
        device_path = self.get_underlying_device(device_path)

        if device_path is None:
            return None

        # Get the available size of the filesystem using `df`
        try:
            result = subprocess.run(
                ["df", "--output=avail", "--block-size=1M", device_path],
                capture_output=True,
                text=True,
                check=True
            )
            # Extract the available size in MiB from the output
            available_mib = float(result.stdout.splitlines()[1].strip())
            return available_mib
        except (subprocess.CalledProcessError, FileNotFoundError, ValueError, IndexError):
            return None

    # ===========================================
    # Data functions
    # ===========================================

    def parse_crypttab(self, crypttab_path: str="/etc/crypttab") -> dict:
        """
        Parse /etc/crypttab and return a dictionary with partition data.

        Args:
            crypttab_path: Path to the crypttab file (default: /etc/crypttab).

        Returns:
            dict: {name: {"device": str, "password": str, "options": str}}
        """
        crypttab = {}
        with open(file=crypttab_path, mode="r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                # Skip comments and empty lines
                if line.startswith("#") or not line:
                    continue
                # Split on whitespace, but not inside quoted strings
                parts = line.split()
                if len(parts) >= 2:
                    name = parts[0]
                    crypttab[name] = {
                        "device": parts[1],
                        "password": parts[2] if len(parts) > 2 else None,
                        "options": parts[3] if len(parts) > 3 else None,
                    }
        return crypttab

    def get_filesystem(self, device_path: str) -> Optional[str]:
        """
        Get the filesystem type of a device or mapped device.
        If the device is a mapped/encrypted device (e.g., /dev/mapper/luks-UUID),
        the filesystem of the underlying device is returned.

        Args:
            device_path: Path to the device (e.g., "/dev/sda1" or "/dev/mapper/luks-UUID").

        Returns:
            Optional[str]: Filesystem type (e.g., "ext4", "vfat") or None if not found.
        """
        # Get the underlying device if it is a mapped device
        device_path = self.get_underlying_device(device_path)

        if device_path is None:
            return None

        # Try to get filesystem with blkid (requires root for most devices)
        if os.geteuid() == 0:
            try:
                result = subprocess.run(
                    ["blkid", "-s", "TYPE", "-o", "value", device_path],
                    capture_output=True,
                    text=True,
                    check=True
                )
                return result.stdout.strip().lower()
            except subprocess.CalledProcessError:
                pass  # Fall through to non-root method

        # Fallback: Check /proc/mounts for mounted devices (no root needed)
        try:
            with open(file="/proc/mounts", mode="r", encoding="utf-8") as f:
                for line in f:
                    parts = line.split()
                    if len(parts) >= 3 and parts[0] == device_path:
                        return parts[2].lower()  # Filesystem type is the 3rd field
        except IOError:
            pass

        # Fallback: Check /dev/disk/by-uuid/ (no root needed)
        try:
            for uuid in os.listdir("/dev/disk/by-uuid"):
                uuid_path = os.path.join("/dev/disk/by-uuid", uuid)
                if os.path.realpath(uuid_path) == device_path:
                    try:
                        result = subprocess.run(
                            ["blkid", "-s", "TYPE", "-o", "value", device_path],
                            capture_output=True,
                            text=True,
                            check=True
                        )
                        return result.stdout.strip().lower()
                    except subprocess.CalledProcessError:
                        pass
        except FileNotFoundError:
            pass

        return None
