#!/sbin/sh

OUTFD="/proc/self/fd/$2"
ZIPFILE="$3"
TARGET_DEVICES=fuxi
READONLY_BLOCKS=
BLOCK_IDS=
STAGED_IMAGE="/tmp/xffc-image-$$"
umask 077

ui_print() {
  printf 'ui_print %s\nui_print\n' "$1" >> "$OUTFD"
}

abort() {
  ui_print "$1"
  exit 1
}

require_tools() {
  for tool in blockdev dd getprop head readlink rm sed sha256sum stat tr unzip; do
    command -v "$tool" >/dev/null 2>&1 || abort "[ERROR] Missing required command: $tool"
  done
  (set -o pipefail) 2>/dev/null || abort "[ERROR] Recovery shell does not support pipefail"
  set -o pipefail

  blockdev_help=$(blockdev --help 2>&1 || :)
  for option in --getsize64 --getro --setrw --setro; do
    case "$blockdev_help" in *"$option"*) ;; *)
      abort "[ERROR] blockdev does not support $option" ;;
    esac
  done

  probe="/tmp/xffc-dd-probe-$$"
  printf x | dd of="$probe" bs=1 conv=fsync >/dev/null 2>&1 || {
    rm -f "$probe"
    abort "[ERROR] dd does not support conv=fsync"
  }
  [ "$(dd if="$probe" iflag=count_bytes count=1 2>/dev/null)" = x ] || {
    rm -f "$probe"
    abort "[ERROR] dd does not support byte-count reads"
  }
  rm -f "$probe"
}

is_target_device() {
  value="$1"
  for target in $TARGET_DEVICES; do
    [ "$value" = "$target" ] && return 0
  done
  return 1
}

is_target_device_list() {
  values=$(printf '%s' "$1" | tr '|,:' '\n\n\n' | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
  while IFS= read -r value; do
    is_target_device "$value" && return 0
  done <<EOF
$values
EOF
  return 1
}

validate_device() {
  for property in ro.product.device ro.product.name ro.product.vendor.device; do
    is_target_device "$(getprop "$property" 2>/dev/null)" && return 0
  done
  is_target_device_list "$(getprop ro.twrp.target.devices 2>/dev/null)" && return 0
  abort "[ERROR] This firmware is only for: $TARGET_DEVICES"
}

read_bootarg() {
  key="$1"
  for source in /proc/cmdline /proc/bootconfig; do
    [ -r "$source" ] || continue
    value=$(
      sed -e 's/ = /=/g' -e 's/"//g' "$source" |
        tr ' ' '\n' |
        sed -n "s/^${key}=//p" |
        head -n 1
    )
    [ -n "$value" ] && printf '%s\n' "$value" && return 0
  done
  return 1
}

set_slots() {
  active=$(getprop ro.boot.slot_suffix 2>/dev/null)
  [ -n "$active" ] || active=$(getprop ro.boot.slot 2>/dev/null)
  [ -n "$active" ] || active=$(read_bootarg androidboot.slot_suffix)
  [ -n "$active" ] || active=$(read_bootarg androidboot.slot)

  case "$active" in
    a|_a) ACTIVE_SLOT=a; INACTIVE_SLOT=b ;;
    b|_b) ACTIVE_SLOT=b; INACTIVE_SLOT=a ;;
    *) abort "[ERROR] Unable to determine active slot" ;;
  esac
}

check_image() {
  partition="$1"
  expected_hash="$2"
  image="firmware-update/$partition.img"
  output=$(unzip -p "$ZIPFILE" "$image" | sha256sum) ||
    abort "[ERROR] Could not read firmware image: $partition"
  actual_hash=${output%% *}
  [ "$actual_hash" = "$expected_hash" ] ||
    abort "[ERROR] Invalid firmware image: $partition"
}

resolve_partition() {
  partition="$1"
  slot="$2"
  for root in \
    /dev/block/by-name \
    /dev/block/bootdevice/by-name \
    /dev/block/platform/*/by-name
  do
    block="$root/${partition}_${slot}"
    if [ -b "$block" ]; then
      readlink -f "$block"
      return
    fi
  done
  return 1
}

restore_readonly() {
  result=0
  for block in $READONLY_BLOCKS; do
    blockdev --setro "$block" >/dev/null 2>&1 || result=1
  done
  return "$result"
}

cleanup() {
  restore_readonly
  rm -f "$STAGED_IMAGE"
}

preflight_partition() {
  partition="$1"
  image_size="$2"

  for slot in a b; do
    block=$(resolve_partition "$partition" "$slot") ||
      abort "[ERROR] Missing block device: ${partition}_${slot}"
    device_id=$(stat -L -c '%t:%T' "$block" 2>/dev/null) ||
      abort "[ERROR] Could not identify block device: ${partition}_${slot}"
    case " $BLOCK_IDS " in
      *" $device_id "*) abort "[ERROR] Duplicate block device: ${partition}_${slot}" ;;
    esac
    BLOCK_IDS="$BLOCK_IDS $device_id"
    eval "BLOCK_${partition}_${slot}=\$block"
    eval "BLOCK_ID_${partition}_${slot}=\$device_id"

    block_size=$(blockdev --getsize64 "$block" 2>/dev/null) ||
      abort "[ERROR] Could not read partition size: ${partition}_${slot}"
    case "$block_size" in
      ''|*[!0-9]*) abort "[ERROR] Invalid partition size: ${partition}_${slot}" ;;
    esac
    [ "$image_size" -le "$block_size" ] ||
      abort "[ERROR] Image is larger than partition: ${partition}_${slot}"

    read_only=$(blockdev --getro "$block" 2>/dev/null) ||
      abort "[ERROR] Could not read partition state: ${partition}_${slot}"
    case "$read_only" in
      0) ;;
      1)
        blockdev --setrw "$block" >/dev/null 2>&1 ||
          abort "[ERROR] Could not make partition writable: ${partition}_${slot}"
        READONLY_BLOCKS="$READONLY_BLOCKS $block"
        [ "$(blockdev --getro "$block" 2>/dev/null)" = 0 ] ||
          abort "[ERROR] Partition is read-only: ${partition}_${slot}"
        ;;
      *) abort "[ERROR] Invalid partition state: ${partition}_${slot}" ;;
    esac
  done
}

flash_partition() {
  partition="$1"
  image_size="$2"
  expected_hash="$3"
  slot="$4"
  image="firmware-update/$partition.img"
  eval "block=\$BLOCK_${partition}_${slot}"
  eval "expected_device_id=\$BLOCK_ID_${partition}_${slot}"

  unzip -p "$ZIPFILE" "$image" > "$STAGED_IMAGE" ||
    abort "[ERROR] Could not stage firmware image: $partition"
  [ "$(stat -c '%s' "$STAGED_IMAGE" 2>/dev/null)" = "$image_size" ] ||
    abort "[ERROR] Staged image has wrong size: $partition"
  output=$(sha256sum "$STAGED_IMAGE") || abort "[ERROR] Could not hash staged image: $partition"
  actual_hash=${output%% *}
  [ "$actual_hash" = "$expected_hash" ] || abort "[ERROR] Invalid staged image: $partition"

  [ -b "$block" ] || abort "[ERROR] Block device disappeared: ${partition}_${slot}"
  [ "$(stat -L -c '%t:%T' "$block" 2>/dev/null)" = "$expected_device_id" ] ||
    abort "[ERROR] Block device changed: ${partition}_${slot}"
  ui_print "[INFO] Updating ${partition}_${slot}..."
  dd if="$STAGED_IMAGE" of="$block" bs=1048576 conv=fsync ||
    abort "[ERROR] Failed to flash ${partition}_${slot}"
  rm -f "$STAGED_IMAGE"

  output=$(dd if="$block" iflag=count_bytes count="$image_size" 2>/dev/null | sha256sum) ||
    abort "[ERROR] Failed to verify ${partition}_${slot}"
  actual_hash=${output%% *}
  [ "$actual_hash" = "$expected_hash" ] ||
    abort "[ERROR] Verification failed: ${partition}_${slot}"
}

flash_slot() {
  slot="$1"
flash_partition abl 233472 fb875d3e15894a0726203f2b706248122b20c8a3aa32dbe2dbd263de1ca42d47 "$slot"
flash_partition aop 282624 d152f6cbc66eddebfda9846edf005e3c80d3c3e6ed052a0f8a55f927a15c5a07 "$slot"
flash_partition aop_config 16384 d8088890a99666e07550d3152d28739792a0c71a676893b37a9698c4aa3e7aea "$slot"
flash_partition bluetooth 1257472 e4f529b33ecac5e9d9bbb4020d5d05423d11809ba77270f4a4de8b6aa5a2036b "$slot"
flash_partition cpucp 192512 9a92523102b46bb693dd013fde044e025280df2b150dfc16ac6c8767a48f4857 "$slot"
flash_partition devcfg 57344 7219cd4ebe465a89b0c8f39b90046cf77e5348a7ea1110b5cd4a66a9cc6faaf6 "$slot"
flash_partition dsp 67108864 e13efaea822c82e4a1a3a8d1958f807a0450e97b9fc3035f04617e4a292cf07a "$slot"
flash_partition featenabler 98304 8e31266a3460955f9ae7ca9b5f18ecc731e1cf96f100e94600cf91d78eb59b8c "$slot"
flash_partition hyp 1458176 7c7707871438d7d3c4f0e578694b7f2740009b3dc3f7c250e401b72440696af1 "$slot"
flash_partition imagefv 6299648 140302e48fb8d7367cf1da5e2c16ea6697df29c700f3e9e9d40cbc9664576123 "$slot"
flash_partition keymaster 409600 ea6013f7989c11aebf5e6cf535ebcb331d58403b05c629779d260ae24d8949c0 "$slot"
flash_partition modem 372338688 48f5443d7874ffe9bb12f68a7c1e35305c8de17466bcc7644b736df3e080fd14 "$slot"
flash_partition multiimgqti 12288 8a3d39667a617b54aaa8944df035e81c6e5b45a3d4fe3db37a0ad6f4271dad11 "$slot"
flash_partition qupfw 57344 24c6219fc6df2ee40059a6ac38e0a5056938cfd9d166e608b73b9b3ba070cf79 "$slot"
flash_partition shrm 65536 a703b0f4b87b8b51093b7ac28ea17de859bab6759c5d1fccab1f0da410849ef7 "$slot"
flash_partition tz 3973120 c62799016e54568992ba603ea7989597aae6deb64dc2b82c99ab6101c6c97f34 "$slot"
flash_partition uefi 2945024 67191e77c9da60bf73c08049ad66113945ad77e512c81c000809c55b8d032747 "$slot"
flash_partition uefisecapp 180224 31320316f48f59f1cda5e8603df814fe6aa4077da3720b9d33c8d39a985ed0fb "$slot"
flash_partition xbl 1200128 8bb5f7bd052ea2785b1c9b3738b26ce68834a1e2d7b98a4b1e30ae5d259cf4f8 "$slot"
flash_partition xbl_config 188416 899e8996df1c9e3f5a433f5d5d6f29cde55d6a5edbdbbfc6f64766dafb99478d "$slot"
flash_partition xbl_ramdump 884736 db709d631383ad0c251ab27cd162b3b14f4b55d25ac21fea666acab7e6ca8911 "$slot"
}

trap cleanup 0
trap 'exit 1' HUP INT TERM

ui_print "####################################################"
ui_print "# Generated by Xiaomi Flashable Firmware Creator"
ui_print "####################################################"
ui_print " "

require_tools
validate_device
set_slots
rm -f "$STAGED_IMAGE"

check_image abl fb875d3e15894a0726203f2b706248122b20c8a3aa32dbe2dbd263de1ca42d47
check_image aop d152f6cbc66eddebfda9846edf005e3c80d3c3e6ed052a0f8a55f927a15c5a07
check_image aop_config d8088890a99666e07550d3152d28739792a0c71a676893b37a9698c4aa3e7aea
check_image bluetooth e4f529b33ecac5e9d9bbb4020d5d05423d11809ba77270f4a4de8b6aa5a2036b
check_image cpucp 9a92523102b46bb693dd013fde044e025280df2b150dfc16ac6c8767a48f4857
check_image devcfg 7219cd4ebe465a89b0c8f39b90046cf77e5348a7ea1110b5cd4a66a9cc6faaf6
check_image dsp e13efaea822c82e4a1a3a8d1958f807a0450e97b9fc3035f04617e4a292cf07a
check_image featenabler 8e31266a3460955f9ae7ca9b5f18ecc731e1cf96f100e94600cf91d78eb59b8c
check_image hyp 7c7707871438d7d3c4f0e578694b7f2740009b3dc3f7c250e401b72440696af1
check_image imagefv 140302e48fb8d7367cf1da5e2c16ea6697df29c700f3e9e9d40cbc9664576123
check_image keymaster ea6013f7989c11aebf5e6cf535ebcb331d58403b05c629779d260ae24d8949c0
check_image modem 48f5443d7874ffe9bb12f68a7c1e35305c8de17466bcc7644b736df3e080fd14
check_image multiimgqti 8a3d39667a617b54aaa8944df035e81c6e5b45a3d4fe3db37a0ad6f4271dad11
check_image qupfw 24c6219fc6df2ee40059a6ac38e0a5056938cfd9d166e608b73b9b3ba070cf79
check_image shrm a703b0f4b87b8b51093b7ac28ea17de859bab6759c5d1fccab1f0da410849ef7
check_image tz c62799016e54568992ba603ea7989597aae6deb64dc2b82c99ab6101c6c97f34
check_image uefi 67191e77c9da60bf73c08049ad66113945ad77e512c81c000809c55b8d032747
check_image uefisecapp 31320316f48f59f1cda5e8603df814fe6aa4077da3720b9d33c8d39a985ed0fb
check_image xbl 8bb5f7bd052ea2785b1c9b3738b26ce68834a1e2d7b98a4b1e30ae5d259cf4f8
check_image xbl_config 899e8996df1c9e3f5a433f5d5d6f29cde55d6a5edbdbbfc6f64766dafb99478d
check_image xbl_ramdump db709d631383ad0c251ab27cd162b3b14f4b55d25ac21fea666acab7e6ca8911

preflight_partition abl 233472
preflight_partition aop 282624
preflight_partition aop_config 16384
preflight_partition bluetooth 1257472
preflight_partition cpucp 192512
preflight_partition devcfg 57344
preflight_partition dsp 67108864
preflight_partition featenabler 98304
preflight_partition hyp 1458176
preflight_partition imagefv 6299648
preflight_partition keymaster 409600
preflight_partition modem 372338688
preflight_partition multiimgqti 12288
preflight_partition qupfw 57344
preflight_partition shrm 65536
preflight_partition tz 3973120
preflight_partition uefi 2945024
preflight_partition uefisecapp 180224
preflight_partition xbl 1200128
preflight_partition xbl_config 188416
preflight_partition xbl_ramdump 884736

ui_print "[INFO] Updating inactive slot $INACTIVE_SLOT first..."
flash_slot "$INACTIVE_SLOT"
ui_print "[INFO] Updating active slot $ACTIVE_SLOT..."
flash_slot "$ACTIVE_SLOT"

restore_readonly || abort "[ERROR] Failed to restore partition read-only state"
READONLY_BLOCKS=
rm -f "$STAGED_IMAGE"
trap - 0
ui_print "Firmware has been successfully updated!"
