# vim: filetype=sh

EMPTY_VOLUME_DATA_SIZE_MB=10
MIN_GROWABLE_FAT_SIZE_MB=256
LVM_PART_OVERHEAD_MB=4      # PV metadata size
DISK_OVERHEAD_MB=1          # offset of 1st partition
BIOSBOOT_PARTITION_SIZE_MB=1
DRAFT_VOLUME_SIZE_ADDUP_MB=$((4*1024))

apply_overhead() {
    size_mb=$1
    vol_subtype=$2
    case "$vol_subtype" in
        ext4)
            # we use a block size of 4K (cf. file 'formatting'), so journal is at least 4M (cf. man mkfs.ext4)
            # this makes an important overhead on small filesystems.
            apply_overhead_percent $((size_mb+8)) 15
            ;;
        fat|efi)
            apply_overhead_percent $((size_mb+4)) 15
            ;;
    esac
}

get_fdisk_type()
{
    part_table_type="$1"
    vol_subtype="$2"
    case "$part_table_type-$vol_subtype" in
        "dos-ext4")
            echo 83
            ;;
        "dos-fat")
            echo c
            ;;
        "dos-lvm")
            echo 8e
            ;;
        "gpt-ext4")
            echo 0FC63DAF-8483-4772-8E79-3D69D8477DE4
            ;;
        "gpt-fat")
            echo EBD0A0A2-B9E5-4433-87C0-68B6B72699C7
            ;;
        "gpt-lvm")
            echo E6D6D379-F507-44C2-A23C-238F2A3DF928
            ;;
        "gpt-bios")
            echo 21686148-6449-6E6F-744E-656564454649
            ;;
        "gpt-efi")
            echo C12A7328-F81F-11D2-BA4B-00A0C93EC93B
            ;;
        *)
            layout_error "partition type '$vol_subtype' is not allowed with a '$part_table_type' partition table"
            ;;
    esac
}

layout_error() {
    echo >&2
    echo "E: invalid disk layout -- $1" >&2
    exit_wrong_user_input
}

save_volume_info() {
    vol_type="$1"
    vol_dir="$2"
    shift 2
    # check number of args
    [ "$4" != "" ] || layout_error "invalid syntax '$*'"
    shift
    mountpoint="$1"
    subtype="$2"
    size="$3"
    # check mountpoint
    [ "$mountpoint" = "none" -o "${mountpoint:0:1}" = "/" ] || \
        layout_error "invalid mount point: '$mountpoint' (should start with '/')"
    # check subtype
    case $subtype in
        fat|ext4)       # ok
            ;;
        efi|bios|lvm)
            [ "$vol_type" = "part" ] || \
                layout_error "'$subtype' type is allowed for partitions, not LVM volumes"
            [ "$subtype" = "efi" -o "$mountpoint" = "none" ] || \
                layout_error "'$subtype' partitions cannot be mounted, use 'none' as a mountpoint"
            ;;
        *)
            [ "$vol_type" = "lvm" ] && \
                layout_error "unknown lvm volume type '$subtype' -- allowed types: fat, ext4"
            [ "$vol_type" = "part" ] && \
                layout_error "unknown partition type '$subtype' -- allowed types: fat, ext4, efi, bios, lvm"
            ;;
    esac
    # check size
    case $size in
        max)       # ok
            ;;
        auto)
            [ "$subtype" = "bios" -o "$mountpoint" != "none" ] || \
                layout_error "cannot set size of '$subtype' partition to 'auto' unless you specify the mountpoint"
            ;;
        *[0-9][%GM])
            check_integer ${size%?}    # last char removed, must be an integer
            ;;
        *)
            layout_error "invalid size '$size'"
            ;;
    esac
    # ok save info
    mkdir -p $vol_dir
    # compute a volume label when relevant
    case $subtype in
        fat|ext4)
            if [ "$mountpoint" = "/" ]
            then
                echo "ROOT_FS" > $vol_dir/label
            elif [ "$mountpoint" != "none" ]
            then
                basename "$mountpoint" | tr a-z A-Z > $vol_dir/label
            elif [ "$vol_type" = "lvm" ]
            then
                # having a label is mandatory for LVM volumes
                echo "lvol$(basename $vol_dir)" > $vol_dir/label
            fi
            ;;
        efi)
            echo EFI > $vol_dir/label
            ;;
    esac
    echo "$vol_type" > $vol_dir/type
    echo "$mountpoint" > $vol_dir/mountpoint
    echo "$subtype" > $vol_dir/subtype
    echo "$size" > $vol_dir/size
}

get_vol_attr_files() {
    attr="$1"
    ls -1 | sort -n | sed -e "s/$/\/$attr/"
}

count_attr_matches() {
    cd "$1"
    files="$(get_vol_attr_files $2)"
    if [ "$files" = "" ]
    then
        echo 0
    else
        grep $3 $files | wc -l
    fi
    cd - >/dev/null
}

check_layout() {
    in_layout_dir=$1

    # check that at least one partition was defined
    [ -d "$in_layout_dir/partitions" ] || \
        layout_error "no partitions defined"

    # check that, except for the last one, size of partitions is not "max" or "<xx>%"
    cd "$in_layout_dir/partitions"
    size_files="$(get_vol_attr_files size | head -n -1)"
    [ $(grep max $size_files | wc -l) -eq 0 ] || \
        layout_error "'max' keyword is only allowed for the last partition"
    [ $(grep "%$" $size_files | wc -l) -eq 0 ] || \
        layout_error "only the last partition can have its size declared as a percentage"
    cd - >/dev/null

    # check that no more than 1 lvm volume has size="max"
    if [ -d "$in_layout_dir/lvm_volumes" ]
    then
        [ $(count_attr_matches "$in_layout_dir/lvm_volumes" size max) -lt 2 ] || \
            layout_error "'max' keyword is only allowed for at most 1 lvm volume"
    fi

    # check that at most one lvm, efi, bios partition is declared
    for special_subtype in lvm efi bios
    do
        [ $(count_attr_matches "$in_layout_dir/partitions" subtype $special_subtype) -lt 2 ] || \
            layout_error "cannot have several partitions with type '$special_subtype'"
    done

    # check if an lvm partition exists
    lvm_partition_exists=$(count_attr_matches "$in_layout_dir/partitions" subtype lvm)

    # if lvm volumes are declared, verify that an lvm partition was declared
    if [ $(ls -1 "$in_layout_dir/lvm_volumes" | wc -l) -gt 0 ]
    then
        [ $lvm_partition_exists -eq 1 ] || \
            layout_error "cannot declare lvm volumes unless you declare a partition with type lvm"
    fi

    # check that mount points are not repeated
    mountpoints="$(cat "$in_layout_dir"/*/*/mountpoint | grep -vx none)"
    [ $(echo "$mountpoints" | wc -w) -eq $(echo "$mountpoints" | uniq | wc -w) ] || \
        layout_error "repeated mount point"

    # check that vg renaming can be handled
    if [ -f "$in_layout_dir"/final_vg_name ]
    then
        [ $lvm_partition_exists -eq 1 ] || \
            layout_error "cannot record LVM volume group name since LVM is not used in this disk layout"
        if ! $target_get_vg_rename_command_exists
        then
            layout_error "this target does not support lvm vg renaming"
        fi
    fi
}

volume_is_partition() {
    parent_dir="$(basename "$(dirname "$1")")"
    [ "$parent_dir" = "partitions" ] && return 0    # yes
    return 1                                        # no
}

check_layout_updates() {
    in_def_layout_dir=$1
    in_new_layout_dir=$2

    # check that important data was not removed or changed from the default layout
    for defvol in $in_def_layout_dir/*/*
    do
        defvol_subtype=$(cat "$defvol/subtype")
        defvol_mountpoint=$(cat "$defvol/mountpoint")

        newvol_found=0
        for newvol in $in_new_layout_dir/*/*
        do
            newvol_subtype=$(cat "$newvol/subtype")
            newvol_mountpoint=$(cat "$newvol/mountpoint")
            if [ "$defvol_mountpoint" = "none" -a "$newvol_subtype" = "$defvol_subtype" ]
            then
                newvol_found=1
                break
            fi
            if [ "$defvol_mountpoint" != "none" -a "$newvol_mountpoint" = "$defvol_mountpoint" ]
            then
                [ "$newvol_subtype" = "$defvol_subtype" ] || \
                    layout_error "'$defvol_mountpoint' volume type cannot be changed from default disk layout ($defvol_subtype)"
                newvol_found=1
                break
            fi
        done

        if [ $newvol_found -eq 0 ]
        then
            [ "$defvol_mountpoint" = "none" ] && info_vol1="$defvol_subtype" || info_vol1="'$defvol_mountpoint'"
            volume_is_partition "$defvol" && info_vol2="partition" || info_vol2="lvm volume"
            layout_error "$info_vol1 $info_vol2 was removed from default disk layout"
        fi
    done

    [ $(cat "$in_def_layout_dir/part_table_type") = $(cat "$in_new_layout_dir/part_table_type") ] || \
        layout_error "changing partition table type (gpt <-> dos) from default disk layout is not allowed"
}

save_lvm_vg_name() {
    in_layout_dir="$1"
    vg_name="$2"
    if [ ! -z "$3" -o -z "$vg_name" ]
    then
        layout_error "lvm vg name declaration should be 'lvm_vg_name <vg_name>' or 'lvm_vg_name auto'"
    fi
    if [ "$vg_name" = "auto" ]
    then
        return  # nothing to do, the auto-generated vg name will be preserved
    fi

    # if string delimiters were used, remove them
    vg_name="$(echo "$vg_name" | tr -d '"'"'")"
    if [ ! -z "$(echo "$vg_name" | tr -d '[a-gA-Z0-9_]')" ]
    then
        layout_error "lvm vg name accepts only the set of chars [a-zA-Z0-9_]"
    fi
    echo "$vg_name" > $in_layout_dir/final_vg_name
}

parse_layout() {
    in_layout_dir="$1"
    in_layout_file="$2"

    mkdir -p "$in_layout_dir/partitions" "$in_layout_dir/lvm_volumes"

    sed -e 's/#.*$//' $in_layout_file | while read inst args
    do
        case "$inst" in
            "")
                ;;
            "partition")
                part_num=$(($(num_dir_entries "$in_layout_dir/partitions")+1))
                part_dir="$in_layout_dir/partitions/$part_num"
                save_volume_info part "$part_dir" "$inst" $args
                ;;
            "lvm_volume")
                vol_num=$(($(num_dir_entries "$in_layout_dir/lvm_volumes")+1))
                vol_dir="$in_layout_dir/lvm_volumes/$vol_num"
                save_volume_info lvm "$vol_dir" "$inst" $args
                ;;
            "gpt"|"dos")
                [ ! -f "$in_layout_dir/part_table_type" ] || \
                    layout_error "several declarations of the partition table type (gpt|dos)"
                echo "$inst" > $in_layout_dir/part_table_type
                ;;
            "lvm_vg_name")
                save_lvm_vg_name "$in_layout_dir" $args
                ;;
            *)
                layout_error "invalid syntax '$inst'"
        esac
    done

    # check that partition table type was defined
    [ -f "$in_layout_dir/part_table_type" ] || \
        layout_error "partition table type is not defined (gpt or dos)"

    # compute fdisk type of partitions
    part_table_type=$(cat "$in_layout_dir"/part_table_type)
    for vol in $(ls -1d "$in_layout_dir"/partitions/*)
    do
        vol_subtype=$(cat $vol/subtype)
        get_fdisk_type $part_table_type $vol_subtype > $vol/fdisk_type
    done
}

estimate_minimal_vol_size_mb()
{
    in_vol="$1"
    in_target_fs="$2"

    vol_mountpoint=$(cat "$in_vol/mountpoint")
    vol_subtype=$(cat "$in_vol/subtype")
    target_fs_path="$in_target_fs/$vol_mountpoint"

    if [ -d "$target_fs_path" ]
    then
        data_size_mb=$(estimated_size_mb "$target_fs_path")
    else
        data_size_mb=$EMPTY_VOLUME_DATA_SIZE_MB
    fi

    min_size_mb=$(apply_overhead $data_size_mb $vol_subtype)

    # if volume is FAT and it should grow during init procedure
    # then ensure size is not too small.
    # otherwise it will be FAT-12 and fatresize will not be able
    # to handle it.
    if ! is_fixed_size_partition $in_vol
    then
        if [ "$vol_subtype" = "fat" -a \
             "$min_size_mb" -lt "$MIN_GROWABLE_FAT_SIZE_MB" ]
        then
            min_size_mb=$MIN_GROWABLE_FAT_SIZE_MB
        fi
    fi
    echo "$min_size_mb"
}

should_extend_up_to_the_end()
{
    vol=$1
    num_last_vol=$(ls -1 "$vol"/.. | sort -n | tail -n 1)
    [ $(basename $vol) = $num_last_vol ]
}

is_fixed_size_partition()
{
    vol="$1"
    vol_type=$(cat "$vol/type")
    if [ "$vol_type" = "lvm" ]
    then
        return 1    # false
    fi
    vol_size=$(cat "$vol/size")
    case "$vol_size" in
        *[GM])
            return 0    # true
            ;;
        *)
            return 1    # false
            ;;
    esac
}

store_computed_size()
{
    mode="$1"
    vol="$2"

    read needed_size_mb

    # in the case of the 'draft' image, we should apply a big margin,
    # except for the partition with subtype "lvm".
    # (the size of this specific partition is obtained by summing the size
    # of lvm volumes, and all of them already include a margin.)
    if [ "$mode" = "draft" -a $(cat "$vol"/subtype) != "lvm" ]
    then
        needed_size_mb=$((needed_size_mb + DRAFT_VOLUME_SIZE_ADDUP_MB))
    fi

    echo $needed_size_mb > $vol/needed_size_mb
}

compute_applied_sizes()
{
    in_mode="$1"
    in_layout_dir="$2"
    in_target_fs="$3"

    # estimate data size of each volume with a mountpoint
    #
    # we have to take care not adding up size of sub-mounts
    # such as in the case of '/' and '/boot' for instance.
    #
    # precedure:
    # 1 - 1st loop: estimate whole tree size at each mountpoint
    #     (ignoring existence of possible sub-mounts)
    # 2 - intermediary line (with 'sort | tac' processing):
    #     sort to get deepest sub-mounts first
    # 3 - 2nd loop: remove from first datasize estimation
    #     the datasize of sub-mounts. To ease this, we use
    #     a temporary file hierarchy "$data_size_tree":
    #     for each mountpoint, we record in parent dirs, up to "/",
    #     the datasize we already counted at this step.
    data_size_tree=$(mktemp -d)
    for vol in $(ls -1d "$in_layout_dir"/*/* | sort -n)
    do
        vol_mountpoint=$(cat "$vol/mountpoint")
        if [ "$vol_mountpoint" != "none" ]
        then
            target_fs_path="$in_target_fs/$vol_mountpoint"
            if [ -d "$target_fs_path" ]
            then
                data_size_mb=$(estimated_size_mb "$target_fs_path")
            else
                data_size_mb=0
            fi
            echo "$vol_mountpoint" "$data_size_mb" "$vol"
        fi
    done | sort -k 1,1 | tac | while read vol_mountpoint data_size_mb vol
    do
        mp=$vol_mountpoint
        mkdir -p "$data_size_tree/$mp"
        if [ -f "$data_size_tree/$mp/.submounts_data_size_mb" ]
        then
            submounts_data_size_mb=$(cat $data_size_tree/$mp/.submounts_data_size_mb)
            data_size_mb=$((data_size_mb-submounts_data_size_mb))
        fi
        echo $data_size_mb > $vol/data_size_mb
        while [ "$mp" != "/" ]
        do
            mp=$(dirname $mp)
            if [ -f "$data_size_tree/$mp/.submounts_data_size_mb" ]
            then
                submounts_data_size_mb=$(cat $data_size_tree/$mp/.submounts_data_size_mb)
            else
                submounts_data_size_mb=0
            fi
            echo $((data_size_mb+submounts_data_size_mb)) > \
                        "$data_size_tree/$mp/.submounts_data_size_mb"
        done
    done
    rm -rf $data_size_tree

    # compute the size we will apply to each volume.
    #
    # we sort volumes to list lvm volumes before partitions.
    # this allows to know the size of all lvm volumes when
    # we have to compute the size of the lvm-type partition.
    for vol in $(ls -1d "$in_layout_dir"/*/* | sort)
    do
        vol_size=$(cat "$vol/size")
        vol_type=$(cat "$vol/type")
        vol_subtype=$(cat "$vol/subtype")
        if [ "$vol_subtype" = "bios" ]
        then
            # bios boot partition is special, we know which size is recommended
            echo $BIOSBOOT_PARTITION_SIZE_MB
        elif [ "$vol_subtype" = "lvm" ]
        then
            # lvm-type partition => sum the size of all lvm volumes, add lvm metadata size
            {   cat "$in_layout_dir"/lvm_volumes/*/needed_size_mb
                echo $LVM_PART_OVERHEAD_MB
            } | sum_lines
        else
            # variable size
            estimate_minimal_vol_size_mb $vol $in_target_fs
        fi | store_computed_size $in_mode $vol
    done

    # re-sort lvm volumes, smallest needed_size first
    # (otherwise approximation may cause the last volume to be short on disk space)
    mkdir "$in_layout_dir"/sorted_lvm_volumes/
    for vol_id in $(ls -1 "$in_layout_dir"/lvm_volumes/)
    do
        vol="$in_layout_dir"/lvm_volumes/$vol_id
        echo "$(cat $vol/needed_size_mb) $vol"
    done | sort -k 1,1 | {
        i=1
        while read needed_size_mb vol
        do
            cp -r $vol "$in_layout_dir"/sorted_lvm_volumes/$i
            i=$((i+1))
        done
    }
    rm -rf "$in_layout_dir"/lvm_volumes
    mv "$in_layout_dir"/sorted_lvm_volumes "$in_layout_dir"/lvm_volumes

    # compute applied_size_mb
    for vol in $(ls -1d "$in_layout_dir"/*/* | sort)
    do
        if should_extend_up_to_the_end $vol
        then
            echo 0 > $vol/applied_size_mb   # 0 means 'fill available space'
        elif is_fixed_size_partition $vol
        then
            size_as_mb $(cat $vol/size) > $vol/applied_size_mb
            # if requested size if too low, set it to minimum needed size
            if [ $(cat $vol/applied_size_mb) -lt $(cat $vol/needed_size_mb) ]
            then
                cp $vol/needed_size_mb $vol/applied_size_mb
            fi
        else
            cp $vol/needed_size_mb $vol/applied_size_mb
        fi
    done

    # compute the overall image size
    {
        for vol in $(ls -1d "$in_layout_dir"/partitions/*)
        do
            if is_fixed_size_partition $vol
            then
                cat $vol/applied_size_mb
            else
                cat $vol/needed_size_mb
            fi
        done
        echo $DISK_OVERHEAD_MB
    } | sum_lines > "$in_layout_dir"/needed_size_mb
}

dump_partition_volumes_info() {
    in_layout_dir="$1"
    for vol in $(ls -1d "$in_layout_dir"/partitions/*)
    do
        echo "$(basename $vol);$(cat $vol/subtype);$(cat $vol/mountpoint);$(cat $vol/size)"
    done
}

dump_lvm_volumes_info() {
    in_layout_dir="$1"
    for vol_id in $(ls -1 "$in_layout_dir"/lvm_volumes/)
    do
        vol="$in_layout_dir"/lvm_volumes/$vol_id
        echo "$(cat $vol/label);$(cat $vol/subtype);$(cat $vol/mountpoint);$(cat $vol/size)"
    done
}

check_need_fatresize() {
    in_layout_dir="$1"
    # last partition may grow
    growable_volumes=$(ls -1d "$in_layout_dir"/partitions/* | tail -n 1)
    # all lvm volumes may grow too
    if [ "$(ls -1 "$in_layout_dir"/lvm_volumes | wc -l)" -gt 0 ]
    then
        growable_volumes="$growable_volumes $(ls -1d "$in_layout_dir"/lvm_volumes/*)"
    fi
    # check if one of them uses a fat filesystem
    for vol in $growable_volumes
    do
        if [ $(cat $vol/subtype) = "fat" ]
        then
            echo 1
            return
        fi
    done
    echo 0
}