# Copyright 2023 Oliver Smith
# SPDX-License-Identifier: GPL-3.0-or-later
from pmb.helpers import logging
import os
from pathlib import Path
from pmb.types import PmbArgs
import pmb.helpers.mount
import pmb.install.losetup
import pmb.helpers.cli
import pmb.config
from pmb.core import Chroot
from pmb.core.context import get_context


def previous_install(path: Path) -> bool:
    """
    Search the disk for possible existence of a previous installation of
    pmOS. We temporarily mount the possible pmOS_boot partition as
    /dev/diskp1 inside the native chroot to check the label from there.
    :param path: path to disk block device (e.g. /dev/mmcblk0)
    """
    label = ""
    for blockdevice_outside in [path.with_stem(f"{path.name}1"), path.with_stem(f"{path.name}p1")]:
        if not os.path.exists(blockdevice_outside):
            continue
        blockdevice_inside = "/dev/diskp1"
        pmb.helpers.mount.bind_file(blockdevice_outside, Chroot.native() / blockdevice_inside)
        try:
            label = pmb.chroot.root(
                ["blkid", "-s", "LABEL", "-o", "value", blockdevice_inside], output_return=True
            )
        except RuntimeError:
            logging.info(
                "WARNING: Could not get block device label,"
                " assume no previous installation on that partition"
            )

        pmb.helpers.run.root(["umount", Chroot.native() / blockdevice_inside])
    return "pmOS_boot" in label


def mount_disk(path: Path) -> None:
    """
    :param path: path to disk block device (e.g. /dev/mmcblk0)
    """
    # Sanity checks
    if not os.path.exists(path):
        raise RuntimeError(f"The disk block device does not exist: {path}")
    for path_mount in path.parent.glob(f"{path.name}*"):
        if pmb.helpers.mount.ismount(path_mount):
            raise RuntimeError(f"{path_mount} is mounted! Will not attempt to format this!")
    logging.info(f"(native) mount /dev/install (host: {path})")
    pmb.helpers.mount.bind_file(path, Chroot.native() / "dev/install")
    if previous_install(path):
        if not pmb.helpers.cli.confirm(
            "WARNING: This device has a previous installation of pmOS. CONTINUE?"
        ):
            raise RuntimeError("Aborted.")
    else:
        if not pmb.helpers.cli.confirm(f"EVERYTHING ON {path} WILL BE ERASED! CONTINUE?"):
            raise RuntimeError("Aborted.")


def create_and_mount_image(
    args: PmbArgs, size_boot: int, size_root: int, size_reserve: int, split: bool = False
) -> None:
    """
    Create a new image file, and mount it as /dev/install.

    :param size_boot: size of the boot partition in MiB
    :param size_root: size of the root partition in MiB
    :param size_reserve: empty partition between root and boot in MiB (pma#463)
    :param split: create separate images for boot and root partitions
    """

    # Short variables for paths
    chroot = Chroot.native()
    config = get_context().config
    img_path_prefix = Path("/home/pmos/rootfs")
    img_path_full = img_path_prefix / f"{config.device}.img"
    img_path_boot = img_path_prefix / f"{config.device}-boot.img"
    img_path_root = img_path_prefix / f"{config.device}-root.img"

    # Umount and delete existing images
    for img_path in [img_path_full, img_path_boot, img_path_root]:
        outside = chroot / img_path
        if os.path.exists(outside):
            pmb.helpers.mount.umount_all(chroot / "mnt")
            pmb.install.losetup.umount(img_path)
            pmb.chroot.root(["rm", img_path])

    # Make sure there is enough free space
    size_mb = round(size_boot + size_reserve + size_root)
    disk_data = os.statvfs(get_context().config.work)
    free = round((disk_data.f_bsize * disk_data.f_bavail) / (1024**2))
    if size_mb > free:
        raise RuntimeError(
            f"Not enough free space to create rootfs image! (free: {free}M, required: {size_mb}M)"
        )

    # Create empty image files
    pmb.chroot.user(["mkdir", "-p", "/home/pmos/rootfs"])
    size_mb_full = str(size_mb) + "M"
    size_mb_boot = str(size_boot) + "M"
    size_mb_root = str(size_root) + "M"
    images = {img_path_full: size_mb_full}
    if split:
        images = {img_path_boot: size_mb_boot, img_path_root: size_mb_root}
    for img_path, image_size_mb in images.items():
        logging.info(f"(native) create {img_path.name} ({image_size_mb})")
        pmb.chroot.root(["truncate", "-s", image_size_mb, img_path])

    # Mount to /dev/install
    mount_image_paths = {img_path_full: "/dev/install"}
    if split:
        mount_image_paths = {img_path_boot: "/dev/installp1", img_path_root: "/dev/installp2"}

    for img_path, mount_point in mount_image_paths.items():
        logging.info(f"(native) mount {mount_point} ({img_path.name})")
        pmb.install.losetup.mount(img_path, args.sector_size)
        device = pmb.install.losetup.device_by_back_file(img_path)
        pmb.helpers.mount.bind_file(device, Chroot.native() / mount_point)


def create(
    args: PmbArgs, size_boot: int, size_root: int, size_reserve: int, split: bool, disk: Path | None
) -> None:
    """
    Create /dev/install (the "install blockdevice").

    :param size_boot: size of the boot partition in MiB
    :param size_root: size of the root partition in MiB
    :param size_reserve: empty partition between root and boot in MiB (pma#463)
    :param split: create separate images for boot and root partitions
    :param disk: path to disk block device (e.g. /dev/mmcblk0) or None
    """
    pmb.helpers.mount.umount_all(Chroot.native() / "dev/install")
    if disk:
        mount_disk(disk)
    else:
        create_and_mount_image(args, size_boot, size_root, size_reserve, split)