#!/bin/bash

# ========= Setup ========= #
CRNT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MOUNT_POINTS=("system" "vendor" "product" "system_ext" "system_dlkm" "vendor_dlkm")

# ========= Colors ========= #
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

pause_and_return() {
    echo
    read -rp $'\033[1;33mPress Enter to return to the menu...\033[0m'
}

# ========= Menu Loop ========= #
while true; do
    clear
    echo -e "${YELLOW}===== A-Team Super Image Partition Mounting Tool =====${NC}"
    echo -e "${CYAN}Current Directory: $CRNT_DIR${NC}"
    echo
    echo "1) Mount Partitions"
    echo "2) Unmount Partitions"
    echo "3) Exit"
    echo
    read -rp "Choose an option [1-3]: " choice

    case "$choice" in
        1)
            echo -e "${CYAN}\nCreating mount folders...${NC}"
            for dir in "${MOUNT_POINTS[@]}"; do
                mkdir -p "$CRNT_DIR/$dir"
            done

            echo -e "${CYAN}Mounting partition images...${NC}"
            sudo mount -o loop "$CRNT_DIR/system.img" "$CRNT_DIR/system/"
            sudo mount -o loop "$CRNT_DIR/vendor.img" "$CRNT_DIR/vendor/"
            sudo mount -o loop "$CRNT_DIR/product.img" "$CRNT_DIR/product/"
            sudo mount -o loop "$CRNT_DIR/system_ext.img" "$CRNT_DIR/system_ext/"
            
            # Handle *_a.img fallback
            if [[ -f "$CRNT_DIR/system_dlkm_a.img" ]]; then
                sudo mount -o loop "$CRNT_DIR/system_dlkm_a.img" "$CRNT_DIR/system_dlkm/"
            elif [[ -f "$CRNT_DIR/system_dlkm.img" ]]; then
                sudo mount -o loop "$CRNT_DIR/system_dlkm.img" "$CRNT_DIR/system_dlkm/"
            fi

            if [[ -f "$CRNT_DIR/vendor_dlkm_a.img" ]]; then
                sudo mount -o loop "$CRNT_DIR/vendor_dlkm_a.img" "$CRNT_DIR/vendor_dlkm/"
            elif [[ -f "$CRNT_DIR/vendor_dlkm.img" ]]; then
                sudo mount -o loop "$CRNT_DIR/vendor_dlkm.img" "$CRNT_DIR/vendor_dlkm/"
            fi

            echo -e "\n${GREEN}✅ All partitions mounted (if available).${NC}"
            pause_and_return
            ;;
        2)
            echo -e "${CYAN}\nUnmounting partitions...${NC}"
            for dir in "${MOUNT_POINTS[@]}"; do
                if mountpoint -q "$CRNT_DIR/$dir"; then
                    sudo umount "$CRNT_DIR/$dir"
                fi
                rm -rf "$CRNT_DIR/$dir"
            done
            echo -e "${GREEN}✅ All unmounted and cleaned up.${NC}"
            pause_and_return
            ;;
        3)
            echo -e "${GREEN}👋 Exiting. See you space cowboy...${NC}"
            break
            ;;
        *)
            echo -e "${RED}❌ Invalid choice. Please enter 1, 2, or 3.${NC}"
            pause_and_return
            ;;
    esac
done