#!/bin/bash

source /etc/tos/scripts/scripts
defaultsys=/etc/.default.sys
MNTYPE="-t ext4"
MNTOPTS=$MNTOPT_EXT4

script_get_lock >/dev/null
[ $? -ne 0 ] && exit 1

usage=$(basename $0):"-a[fix|add|up] -r[/dev/md0] -b[/dev/sda,/dev/sdb]\n"

TOS_ACTION=""
TOS_RAID=""
TOS_DEVICES=""
while getopts "a:r:b:" opt; do
    case $opt in
    a) TOS_ACTION=$OPTARG ;; # get action
    r) TOS_RAID=$OPTARG ;;   # open raid 1 bitmap
    b) TOS_DEVICES=$OPTARG ;;
    esac
done

if [ -z "${TOS_ACTION}" -o -z "${TOS_RAID}" -o -z "${TOS_DEVICES}" ]; then
    echo -e $usage
    script_put_lock >/dev/null
    exit 1
fi

mdopslog="/tmp/mdops.log"
rm -f $mdopslog # remove old tmp file

isblockdev() {
    local blk=$1
    if [ ! -b $blk ]; then
        echo "Error: $blk is not a block device" >> $mdopslog
        return 1
    fi
    return 0
}

ismddev() {
    local raid=$1
    local fc=$(mdadm -D ${raid} | grep ${TOS_FC_SIGN})
    local core=$(mdadm -D ${raid} | grep ${TOS_CORE_SIGN})
    local swap=$(mdadm -D ${raid} | grep ${TOS_SWAP_SIGN})
    local user=$(mdadm -D ${raid} | grep ${TOS_USER_SIGN})
    if [ ! -z "${core}" -o ! -z "${swap}" -o ! -z "${user}" -o ! -z "${fc}" ]; then
        return 0
    else
        return 1
    fi
}

ismdmember() {
    local blk=$1
    local number=$(ls ${blk}? | wc -l)
    if [ $number -eq 4 ]; then
        [ ! -b "${blk}2" -o ! -b "${blk}3" ] && return 1
        mdadm -E "${blk}2" | grep ${TOS_CORE_SIGN}
        [ $? -ne 0 ] && return 1
        mdadm -E "${blk}3" | grep ${TOS_SWAP_SIGN}
        [ $? -ne 0 ] && return 1
        return 0
    else
        return 1
    fi
}

RemoveFromRaid() {
    local blk=$1
    local blkname=$(basename $blk)
    cat /proc/mdstat | grep "$blkname" | while read line; do
        mdname=$(echo $line | awk -F' : ' '{print $1}')
        if [ "$mdname" = "$(basename ${TOS_RAID_CORE})" ]; then
            mdadm --manage $TOS_RAID_CORE --fail ${blk}2 >> $mdopslog  2>&1
            mdadm --manage $TOS_RAID_CORE --remove ${blk}2 >> $mdopslog  2>&1
        elif [ "$mdname" = "$(basename $TOS_RAID_SWAP)" ]; then
            mdadm --manage $TOS_RAID_SWAP --fail ${blk}3 >> $mdopslog  2>&1
            mdadm --manage $TOS_RAID_SWAP --remove ${blk}3 >> $mdopslog  2>&1
        else
            mdadm --manage /dev/$mdname --fail ${blk}4 >> $mdopslog  2>&1
            mdadm --manage /dev/$mdname --remove ${blk}4 >> $mdopslog  2>&1
        fi
    done
}

RemoveFromRaid_4th_partition() {
    local blk=$1
    local blkname=$(basename $blk)
    cat /proc/mdstat | grep "$blkname" | while read line; do
        mdname=$(echo $line | awk -F' : ' '{print $1}')
        if [ "$mdname" != "$(basename ${TOS_RAID_CORE})" ] && [ "$mdname" != "$(basename ${TOS_RAID_SWAP})" ]; then
            mdadm --manage /dev/$mdname --fail ${blk}4 >> $mdopslog  2>&1
            mdadm --manage /dev/$mdname --remove ${blk}4 >> $mdopslog  2>&1
        fi
    done
}

clear_partitions() {
    local blk=$1

    # Get all the partition Numbers and arrange in descending order
    local partitions=$(parted -m "$blk" unit MB print | awk -F: '/^[0-9]+:/ {print $1}' | sort -nr)

    # Collect the partition number (>= 4) that you need to delete.
    local to_delete=()
    for part in $partitions; do
        if [ "$part" -ge 4 ]; then
            to_delete+=("$part")
        fi
    done

    for part in "${to_delete[@]}"; do
        parted --script "$blk" rm "$part"
    done
}

isCoreSwapMember() {
    local blk=$1
    if [ -b "${blk}2" ]; then
        mdadm --detail $TOS_RAID_CORE 2>/dev/null | grep -q "${blk}2" && return 0
    fi
    if [ -b "${blk}3" ]; then
        mdadm --detail $TOS_RAID_SWAP 2>/dev/null | grep -q "${blk}3" && return 0
    fi
    return 1
}

blkparted_partial() {
    local blk=$1
    local bootSize=$TOS_BOOT_SIZE
    local coreSize=$TOS_CORE_SIZE
    local swapSize=$TOS_SWAP_SIZE
    local offset=0
    local offend=0
    clear_partitions $blk
    let offend=$offend+$bootSize
    if [ ! -b ${blk}1 ];then
        parted -s $blk mkpart primary ext2 0% $offend
    fi
    let offset=$offend
    let offend=$offset+$coreSize
    if [ ! -b ${blk}2 ];then
        parted -s $blk mkpart primary ext2 $offset $offend
    fi
    let offset=$offend
    let offend=$offset+$swapSize
    if [ ! -b ${blk}3 ];then
        parted -s $blk mkpart primary ext2 $offset $offend
    fi
    partprobe "${blk}" >> $mdopslog  2>&1
    udevadm settle --timeout=10
    if [ ! -b ${blk}1 ] || [ ! -b ${blk}2 ] || [ ! -b ${blk}3 ];then
        sleep 2
        if [ ! -b ${blk}1 ] || [ ! -b ${blk}2 ] || [ ! -b ${blk}3 ];then
            echo "blkparted_partial Error: $blk can't get valid part" >> $mdopslog
            script_put_lock >/dev/null
            exit 1
        fi
    fi
    let offset=$offend
    parted -s $blk mkpart primary ext2 $offset 100%
    if [ $? -ne 0 ];then
        echo "blkparted_partial Error: $blk can't parted user partition" >> $mdopslog
        script_put_lock >/dev/null
        exit 1
    fi
    partprobe "${blk}" >> $mdopslog  2>&1
    udevadm settle --timeout=10
    sleep 1
    echo "Info: mdadm --zero-superblock clear ${blk}4 metadata" >> $mdopslog
    mdadm --zero-superblock "${blk}4" >> $mdopslog  2>&1
}

blkparted() {
    local blk=$1
    local offset=0
    local offend=0

    parted -s $blk "mktable gpt" >> $mdopslog  2>&1
    if [ $? -ne 0 ]; then
        script_put_lock >/dev/null
        exit 1
    fi

    local bootSize=$TOS_BOOT_SIZE
    local coreSize=$TOS_CORE_SIZE
    local swapSize=$TOS_SWAP_SIZE


    if [ $isOldPart -eq 0 ]; then
      coreSize=2048
      swapSize=1024
    fi
    let offend=$offset+$bootSize
    parted -s $blk mkpart primary ext2 0% $offend
    let offset=$offend
    let offend=$offset+$coreSize
    parted -s $blk mkpart primary ext2 $offset $offend
    let offset=$offend
    let offend=$offset+$swapSize
    parted -s $blk mkpart primary ext2 $offset $offend

    let offset=$offend
    parted -s $blk mkpart primary ext2 $offset 100%
#    parted -s $blk unit s mkpart primary $START $END
    partprobe "${blk}" >> $mdopslog  2>&1
    udevadm settle --timeout=10
    sleep 1
    echo "Info: mdadm --zero-superblock clear ${blk}2、${blk}3 metadata" >> $mdopslog
    mdadm --zero-superblock "${blk}2" >> $mdopslog  2>&1
    mdadm --zero-superblock "${blk}3" >> $mdopslog  2>&1
}

base_raid_check() {
    local raid=$1
    local partition=$2
    if [ -b $raid -a -b "${partition}" ]; then
        local insert=0
        local mdname=$(basename $raid)
        for i in $(ls /sys/block/$mdname/slaves/); do
            local blk=/dev/$i
            [ ! -b $blk ] && mdadm --manage $raid --remove $i
            [ "$blk" = "$partition" ] && insert=1
        done
        if [ $insert -eq 0 ]; then
            if [ "${raid}" != "/dev/md8" -a "${raid}" != "/dev/md9" ]; then
              mdadm --grow ${raid} --raid-devices=$totaldisks --force >> $mdopslog  2>&1
              mdadm -a ${raid} ${partition} >> $mdopslog  2>&1
            fi
            # mdadm --grow ${raid} --raid-devices=$TOS_SYS_NUMBER --force
            # mdadm -a ${raid} ${partition} >> $mdopslog  2>&1
        fi
    fi
}

base_raid_fix() {
    local blk=$1
    # if [ -b "${blk}1" ]; then
    #     e2label ${blk}1 | grep ${TOS_DISK_SIGN} >/dev/null
    #     [ $? -ne 0 ] && /etc/tos/scripts/initboot "${blk}"
    # fi

    base_raid_check "$TOS_RAID_CORE" "${blk}2"
    base_raid_check "$TOS_RAID_SWAP" "${blk}3"
}

main_raid_fix() {
    local blk=$1
    if [ -b "$blk" ]; then
        local mdname=$(basename "$blk")
        local cache_size="/sys/block/$mdname/md/stripe_cache_size"
        [ -e $cache_size ] && echo $TOS_TRIP_SIZE >$cache_size
        for ele in $(ls -1d /sys/block/${mdname}/slaves/*); do
            local elename=$(basename $ele)
            mdadm -r $blk $elename >> $mdopslog  2>&1
        done
        mdadm --detail $blk | awk '/faulty/ {print $6}' | while read line; do
            [ -z "$line" ] && continue
            local elename=$(basename $line)
            mdadm -r $blk $elename >> $mdopslog 2>&1
        done
    fi
}

clear_disk_part() {
  local blk=$1
  for ((i=1;i<=100; i++))
    do
    dev="${blk}"$i
    if [ -b $dev ];then
        if [ $i -eq 2 ];then
          local info=$(mdadm -D /dev/md9 | grep $blk)
          if [ -n "$info" ];then
            mdadm /dev/md9 -f $dev
            mdadm /dev/md9 -r $dev
          fi
        fi

        if [ $i -eq 3 ];then
          local mdinfo=$(mdadm -D /dev/md9 | grep $blk)
          if [ -n "$mdinfo" ];then
            mdadm /dev/md8 -f $dev
            mdadm /dev/md8 -r $dev
          fi
        fi
        echo 'Ignore' | parted $blk rm $i
    fi
  done
}

ismddev $TOS_RAID
if [ $? -ne 0 ]; then
    echo "Error: $TOS_RAID is not a raid device" >> $mdopslog
    script_put_lock >/dev/null
    exit 1
fi

# handle...
main_raid_fix $TOS_RAID
TOS_RAID_NAME=$(basename $TOS_RAID)
TOS_SYNC_STATE=$(cat /sys/block/${TOS_RAID_NAME}/md/sync_action 2>/dev/null)
TOS_RAID_DISKS=$(cat /sys/block/${TOS_RAID_NAME}/md/raid_disks)
TOS_RAID_LEVEL=$(cat /sys/block/${TOS_RAID_NAME}/md/level)

disks=(${TOS_DEVICES//,/ })
userdisks=()
index=0
FC_SIGN=$(mdadm -D ${TOS_RAID} | grep ${TOS_FC_SIGN})

/usr/sbin/ter_partsize -r $TOS_RAID
isOldPart=$?
for blk in ${disks[@]}; do
    isblockdev $blk
    if [ $? -ne 0 ]; then
        script_put_lock >/dev/null
        exit 1
    fi
    if [ "X$FC_SIGN" != "X" ]; then
        RemoveFromRaid $blk
        dd if=/dev/zero of=$blk bs=512K count=1 >/dev/null 2>&1
        parted -s $blk mktable gpt
        parted -s $blk mkpart primary ext2 0% 100%
        userdisks[$index]=${blk}1
        let index=index+1
        continue
    fi
    # check the blk part
    # check $blk if is a md member
    [ $isOldPart -eq 0 ] && clear_disk_part $blk

    removeFlag="/tmp/remove_flag"
    touch $removeFlag
    isCoreSwapMember $blk
    if [ $? -eq 0 ]; then
        RemoveFromRaid_4th_partition $blk
        blkparted_partial $blk
    else
        RemoveFromRaid $blk
        blkparted $blk
    fi

    # Avoiding removing old partitions (such as sdb4) triggers rules under udev,
    rm $removeFlag
    # fix the part of system
    base_raid_fix $blk
    # sum the part of user
    userdisks[$index]=${blk}4
    let index=index+1
done

case $TOS_ACTION in
fix)
    if [ "$TOS_RAID_LEVEL" = "raid0" -o "$TOS_RAID_LEVEL" = "linear" ]; then
        script_put_lock >/dev/null
        exit 1
    fi
    if [ "$TOS_SYNC_STATE" = "idle" ]; then
        totaldisks=$(mdadm --detail $TOS_RAID | awk -F' : ' '/Total Devices/ {print $2}')
        if [ $TOS_RAID_DISKS -gt $totaldisks ]; then
            mdadm --manage $TOS_RAID --add ${userdisks[@]} >> $mdopslog  2>&1
            totaldisks=$(mdadm --detail $TOS_RAID | awk -F' : ' '/Total Devices/ {print $2}')
            if [ "$TOS_RAID" != "/dev/md8" -a "$TOS_RAID" != "/dev/md9" ]; then
              mdadm --grow $TOS_RAID --raid-devices=$totaldisks --force >> $mdopslog  2>&1
            fi
            # mdadm --grow $TOS_RAID --raid-devices=$totaldisks
            mdadm -Ds >/etc/mdadm.conf
            echo "$TOS_ACTION $TOS_RAID success" >> $mdopslog
        fi
    fi
    ;;
add)
    fs=$(blkid -o export $TOS_RAID | awk -F= '/^TYPE=/ {print $2}')
    totaldisks=$(mdadm --detail $TOS_RAID | awk -F' : ' '/Total Devices/ {print $2}')
    if [ "$TOS_RAID_LEVEL" = "linear" ]; then
        if [ "$TOS_RAID" != "/dev/md8" -a "$TOS_RAID" != "/dev/md9" ]; then
          mdadm --grow $TOS_RAID --add ${userdisks[@]} >> $mdopslog  2>&1
        fi
        # mdadm --grow $TOS_RAID --add ${userdisks[@]}
        if [ "$fs" = "LVM2_member" ]; then
            pvresize -y $TOS_RAID
        fi
    elif [ "$TOS_RAID_LEVEL" = "raid1" ]; then
        let totaldisks=totaldisks+${#userdisks[@]}
        mdadm --manage $TOS_RAID --add ${userdisks[0]} >> $mdopslog  2>&1
        if [ "$TOS_RAID" != "/dev/md8" -a "$TOS_RAID" != "/dev/md9" ]; then
          mdadm --grow $TOS_RAID --raid-devices=$totaldisks --force >> $mdopslog  2>&1
        fi
        # mdadm --grow $TOS_RAID --add ${userdisks[@]} --raid-devices=$totaldisks
    elif [ "$TOS_RAID_LEVEL" = "raid5" -o "$TOS_RAID_LEVEL" = "raid6" ]; then
        if [ $TOS_RAID_DISKS -lt $totaldisks ]; then
            if [ "$TOS_RAID" != "/dev/md8" -a "$TOS_RAID" != "/dev/md9" ]; then
              mdadm --grow $TOS_RAID --raid-devices=$totaldisks --force >> $mdopslog  2>&1
            fi
            # mdadm --grow $TOS_RAID --raid-devices=$totaldisks
            echo 1 > /dev/null
        elif [ $TOS_RAID_DISKS -eq $totaldisks ]; then
            let totaldisks=totaldisks+1
            mdadm --manage $TOS_RAID --add ${userdisks[0]} >> $mdopslog  2>&1
            if [ "$TOS_RAID" != "/dev/md8" -a "$TOS_RAID" != "/dev/md9" ]; then
              mdadm --grow $TOS_RAID --raid-devices=$totaldisks --force >> $mdopslog  2>&1
            fi
            # mdadm --grow $TOS_RAID --add ${userdisks[0]} --raid-devices=$totaldisks
        fi
        # optimization resync speed for raid5 & raid6...
        cache_size="/sys/block/$TOS_RAID_NAME/md/stripe_cache_size"
        [ -e $cache_size ] && echo $TOS_TRIP_SIZE >$cache_size
    elif [ "$TOS_RAID_LEVEL" = "raid10" ]; then
        if [ $TOS_RAID_DISKS -lt $totaldisks ]; then
            # mdadm --grow $TOS_RAID --raid-devices=$totaldisks
            if [ "$TOS_RAID" != "/dev/md8" -a "$TOS_RAID" != "/dev/md9" ]; then
              mdadm --grow $TOS_RAID --raid-devices=$totaldisks --force >> $mdopslog  2>&1
            fi
            echo 1 > /dev/null
        elif [ $TOS_RAID_DISKS -eq $totaldisks -a ${#userdisks[@]} = 2 ]; then
            let totaldisks=totaldisks+2
            mdadm --manage $TOS_RAID --add ${userdisks[0]} >> $mdopslog  2>&1
            if [ "$TOS_RAID" != "/dev/md8" -a "$TOS_RAID" != "/dev/md9" ]; then
              mdadm --grow $TOS_RAID --raid-devices=$totaldisks --force >> $mdopslog  2>&1
            fi
            # mdadm --grow $TOS_RAID --add ${userdisks[@]} --raid-devices=$totaldisks
        fi
    fi
    ;;
up)
    if [ "$TOS_RAID_LEVEL" = "raid1" -a $TOS_RAID_DISKS -eq 1 ]; then
        mdadm --manage $TOS_RAID --add ${userdisks[0]} >> $mdopslog  2>&1
        if [ "$TOS_RAID" != "/dev/md8" -a "$TOS_RAID" != "/dev/md9" ]; then
          mdadm --grow $TOS_RAID --raid-devices=2 --force >> $mdopslog  2>&1
        fi
        # mdadm --grow $TOS_RAID --raid-devices=2 --force
    elif [ "$TOS_RAID_LEVEL" = "raid1" -a $TOS_RAID_DISKS -ge 2 ]; then
        let TOS_RAID_DISKS=TOS_RAID_DISKS+1
        mdadm --manage $TOS_RAID --add ${userdisks[0]} >> $mdopslog  2>&1
        if [ "$TOS_RAID" != "/dev/md8" -a "$TOS_RAID" != "/dev/md9" ]; then
          mdadm --grow $TOS_RAID --level=raid5 --raid-devices=$TOS_RAID_DISKS --force >> $mdopslog  2>&1
        fi
        # mdadm --grow $TOS_RAID --level=raid5 --add ${userdisks[0]} --raid-devices=$TOS_RAID_DISKS --force
    elif [ "$TOS_RAID_LEVEL" = "raid5" -a $TOS_RAID_DISKS -ge 3 ]; then
        let TOS_RAID_DISKS=TOS_RAID_DISKS+1
        mdadm --manage $TOS_RAID --add ${userdisks[0]} >> $mdopslog  2>&1
        if [ "$TOS_RAID" != "/dev/md8" -a "$TOS_RAID" != "/dev/md9" ]; then
          mdadm --grow $TOS_RAID --level=raid6 --raid-devices=$TOS_RAID_DISKS --force >> $mdopslog  2>&1
        fi
        # mdadm --grow $TOS_RAID --level=raid6 --add ${userdisks[0]} --raid-devices=$TOS_RAID_DISKS --force
    fi
    # optimization resync speed for raid5 & raid6...
    cache_size="/sys/block/$TOS_RAID_NAME/md/stripe_cache_size"
    [ -e $cache_size ] && echo $TOS_TRIP_SIZE >$cache_size
    mdadm -Ds >/etc/mdadm.conf
    ;;
*)
    echo -e $usage
    ;;
esac

script_put_lock >/dev/null
exit 0
