import logging
import re
import subprocess
from pathlib import Path
from typing import Dict

from .infoenums import PartitionInfo, FsInfo
from .settings import (
    EXPECTED_NUMBER_OF_PARTITIONS,
    PARTITION_LAYOUT,
    OTA_RESCUE_LOGGER_NAME,
    PERSISTENT_DATA_PARTITION,
    PERSISTENT_DATA_PARTITION_START,
    PARTITION_BOOT_ID,
    PARTITION_ROOTFS0_ID,
    PARTITION_ROOTFS1_ID,
)

otalogger = logging.getLogger(OTA_RESCUE_LOGGER_NAME)


class SdCard:
    def __init__(self, deviceName):
        self._deviceName = deviceName
        self._devicePath = "/dev/{}".format(self._deviceName)

    def generatePartitions(self):
        partitions = self._getPartitionLayout()
        fileSystemInformation = self.getFileSystemOverview()

        if EXPECTED_NUMBER_OF_PARTITIONS not in partitions.keys():
            otalogger.error(
                "Expected partition number {} is not available.".format(
                    EXPECTED_NUMBER_OF_PARTITIONS
                )
            )
            return False

        if (
            partitions[EXPECTED_NUMBER_OF_PARTITIONS][PartitionInfo.LABEL]
            != PERSISTENT_DATA_PARTITION
        ):
            otalogger.error(
                "Expected partition label '{}' is not available.".format(
                    PERSISTENT_DATA_PARTITION
                )
            )
            return False

        partitionStartOffsetMiB = int(
            partitions[EXPECTED_NUMBER_OF_PARTITIONS][PartitionInfo.START][:-3]
        )
        if partitionStartOffsetMiB != PERSISTENT_DATA_PARTITION_START:
            otalogger.error(
                "Partition '{}' does not start at the expected position ({}MiB/{}MiB).".format(
                    PERSISTENT_DATA_PARTITION,
                    PERSISTENT_DATA_PARTITION_START,
                    partitionStartOffsetMiB,
                )
            )
            return False

        for partitionNumber in range(1, EXPECTED_NUMBER_OF_PARTITIONS):
            if f"{self._devicePath}p{partitionNumber}" in fileSystemInformation.keys():
                if not self._umountFileSystem(partitionNumber):
                    otalogger.error(
                        f"Unable to unmount partition {self._devicePath}p{partitionNumber}."
                    )
                    return False
            else:
                otalogger.debug(
                    f"Partition number {partitionNumber} is not mounted. Skipping unmounting it."
                )

        for partitionNumber in range(1, EXPECTED_NUMBER_OF_PARTITIONS):
            if partitionNumber in partitions.keys():
                if not self._deletePartition(partitionNumber):
                    otalogger.error(
                        f"Unable to delete partition {self._devicePath}p{partitionNumber}."
                    )
                    return False
            else:
                otalogger.debug(
                    f"Partition number {partitionNumber} does not exist. Skipping deleting it."
                )

        self._reloadPartitionTable()

        for partitionNumber, partition in PARTITION_LAYOUT.items():
            if not self._createPartition(
                partition[PartitionInfo.START], partition[PartitionInfo.END]
            ):
                otalogger.error(
                    f"Unable to create partition `{partition[PartitionInfo.LABEL]}` at position {partitionNumber}."
                )
                return False

        self._reloadPartitionTable()

        bootPartitionInfo = PARTITION_LAYOUT[PARTITION_BOOT_ID]
        self._createBootFileSystem(
            PARTITION_BOOT_ID, bootPartitionInfo[PartitionInfo.LABEL]
        )

        rootFsPartitions = {
            i: PARTITION_LAYOUT[i] for i in (PARTITION_ROOTFS0_ID, PARTITION_ROOTFS1_ID)
        }
        for partitionNumber, partition in rootFsPartitions.items():
            if not self._createFileSystem(
                partitionNumber, partition[PartitionInfo.LABEL]
            ):
                otalogger.error(
                    f"Unable to create ext4 file system on partition {partitionNumber}."
                )
                return False

        self._reloadPartitionTable()

        return True

    def printOverview(self):
        """
        Prints the systems partition overview.
        """
        partitions = self._getPartitionLayout()
        fileSystemInformation = self.getFileSystemOverview()

        for partitionNumber, partition in partitions.items():
            print(
                f"{partitionNumber}:\n\t"
                f"start: {partition[PartitionInfo.START]}\n\t"
                f"end: {partition[PartitionInfo.END]}\n\t"
                f"size: {partition[PartitionInfo.SIZE]}\n\t"
                f"label: {partition[PartitionInfo.LABEL]}"
            )

        for fsName, fs in fileSystemInformation.items():
            print(
                f"{fsName}:\n\t"
                f"size: {fs[FsInfo.TOTAL]}\n\t"
                f"mount point: {fs[FsInfo.MOUNT_DIR]}"
            )

    def mountPartition(self, partitionLabel: str, mountPoint: Path) -> bool:
        """
        Mount partition with the given label at the given mount point.

        :param partitionLabel: The label of the partition to mount.
        :param mountPoint: The mount point.
        :return: Boolean indicating if the partition could successfully be mounted.
        """
        partitionLayout = self._getPartitionLayout()

        selectedPartitionNumber = None
        for partitionNumber, partitionInfo in partitionLayout.items():
            if partitionInfo.get(PartitionInfo.LABEL) == partitionLabel:
                selectedPartitionNumber = partitionNumber

        if selectedPartitionNumber is None:
            otalogger.error(f"No partition with the label `{partitionLabel}` found.")
            return False

        if not self._mountFileSystem(selectedPartitionNumber, mountPoint.as_posix()):
            otalogger.error(
                f"Failed to mount partition with the label `{partitionLabel}`."
            )
            return False

        return True

    def _getPartitionLayout(self) -> Dict:
        """
        Gather the partition layout of a block device via `parted`.

        :return: Dictionary containing a list of all partitions and their properties.
        """
        result = subprocess.run(
            ["parted", self._devicePath, "unit", "MiB", "print"],
            capture_output=True,
            text=True,
        )
        partitions = {}

        if result.returncode != 0:
            otalogger.error("Failed to read partition table.")
            return partitions

        partitionRegex = re.compile(
            r"(\d+)\s+([\d.]+[kMG]iB)\s+([\d.]+[kMG]iB)\s+([\d.]+[kMG]iB)\s+\w+(?:\s+)?(\w+)?(?:\s+[A-z,\s]+)?"
        )

        partitions = {
            int(match.group(1)): {
                PartitionInfo.START: match.group(2),
                PartitionInfo.END: match.group(3),
                PartitionInfo.SIZE: match.group(4),
                PartitionInfo.FILE_SYSTEM: match.group(5),
            }
            for line in result.stdout.strip().splitlines()
            if (match := partitionRegex.search(line))
        }

        def extractPartitionLabel(line):
            match = re.search(r'LABEL="([^"]+)"', line)
            if match:
                return match.group(1)
            return ""

        for partitionNumber, partition in partitions.items():
            result = subprocess.run(
                ["blkid", f"{self._devicePath}p{partitionNumber}"],
                capture_output=True,
                text=True,
            )
            partitions[partitionNumber][PartitionInfo.LABEL] = extractPartitionLabel(
                result.stdout
            )

        return partitions

    def _createPartition(self, start: int, end: int) -> bool:
        """
        Create a new partition.

        :param start: Start offset in megabytes (MiB).
        :param end: End offset in megabytes (MiB).
        :return: Boolean indicating if the partition could be created.
        """
        otalogger.debug(
            f"Creating partition on {self._devicePath} from {start}MiB to {end}MiB."
        )
        result = subprocess.run(
            [
                "parted",
                self._devicePath,
                "mkpart",
                "primary",
                f"{start}MiB",
                f"{end}MiB",
            ],
            stdout=subprocess.DEVNULL,
        )
        return result.returncode == 0

    def _deletePartition(self, partitionNumber: int) -> bool:
        """
        Delete partition with the given partition number.

        :param partitionNumber: The partition number to delete.
        :return: Boolean indicating if the partition could be deleted.
        """
        otalogger.debug(f"Deleting partition {partitionNumber} on {self._devicePath}.")
        result = subprocess.run(
            ["parted", self._devicePath, "rm", f"{partitionNumber}"],
            stdout=subprocess.DEVNULL,
        )
        return result.returncode == 0

    def _reloadPartitionTable(self) -> bool:
        """
        Reloads the partition table without rebooting.

        :return: Boolean indicating if the partition table could be reloaded.
        """
        result = subprocess.run(
            [
                "partprobe",
                self._devicePath,
            ],
            stdout=subprocess.DEVNULL,
        )
        return result.returncode == 0

    def _createBootFileSystem(self, partitionNumber: int, partitionLabel="boot"):
        """
        Creates a files system for the boot partition that requires some special parameters
        for block size and omitting metadata checksum

        :param partitionNumber: partition number on which the file system will be created
        :param partitionLabel: label for the file system to be created
        """
        result = subprocess.run(
            [
                "mkfs.ext4",
                f"{self._devicePath}p{partitionNumber}",
                "-L",
                partitionLabel,
                "-b",
                "1024",
                "-O",
                "^metadata_csum",
            ],
            stdout=subprocess.DEVNULL,
        )
        return result.returncode == 0

    def _createFileSystem(self, partitionNumber: int, partitionLabel: str) -> bool:
        """
        Create a file system on the given partition.

        :param partitionNumber: The partition number.
        :param partitionLabel: The partition's new label.
        :return: Boolean indicating if the file system could be created.
        """
        result = subprocess.run(
            [
                "mkfs.ext4",
                f"{self._devicePath}p{partitionNumber}",
                "-L",
                partitionLabel,
            ],
            stdout=subprocess.DEVNULL,
        )
        return result.returncode == 0

    def _mountFileSystem(self, partitionNumber: int, mountPath: str) -> bool:
        """
        Mount a partition on a specific mount point.

        :param partitionNumber: The partition number.
        :param mountPath: The mount path.
        :return: Boolean indicating if the partition could be mounted.
        """
        result = subprocess.run(
            ["mount", f"{self._devicePath}p{partitionNumber}", mountPath],
            stdout=subprocess.DEVNULL,
        )
        return result.returncode == 0

    def _umountFileSystem(self, partitionNumber: int) -> bool:
        """
        Unmount a partition.

        :param partitionNumber: The partition number.
        :return: Boolean indicating if the partition could be unmounted.
        """
        result = subprocess.run(
            ["umount", f"{self._devicePath}p{partitionNumber}"],
            stdout=subprocess.DEVNULL,
        )
        return result.returncode == 0

    @staticmethod
    def getFileSystemOverview() -> Dict:
        """
        Creates a dict containing information about all file systems found with `df`.

        :return: Dictionary containing a list of all file systems and their properties.
        """
        output = subprocess.check_output(["df", "-hT"])
        all_lines = output.decode("utf-8")
        lines = all_lines.split("\n")

        fsOverview = {}

        for line in lines:
            # skip initial and final (empty) line
            if "filesystem" in line.lower() or len(line) == 0:
                continue

            # format lines to be separated by a single space
            data = re.sub(" +", " ", line)
            fsName, fsType, szTotal, szUsed, szAvailable, usedPercentage, mountPoint = (
                data.split(" ")
            )

            if fsType in ["tmpfs", "rootfs", "devtmpfs"]:
                # skip `tmpfs`, `rootfs` and `devtmpfs` in this list
                continue

            fsOverview[fsName] = {
                FsInfo.TYPE: fsType,
                FsInfo.TOTAL: szTotal,
                FsInfo.USED: szUsed,
                FsInfo.AVAILABLE: szAvailable,
                FsInfo.MOUNT_DIR: mountPoint,
            }

        return fsOverview
