#!/bin/bash
# Script to patch and sign the reMarkable Android APKs

set -euo pipefail

PROJECT_DIR=$(cd "$(dirname "$0")/.." && pwd)
BUILD_DIR="$PROJECT_DIR/build"
OUTPUT_DIR="$PROJECT_DIR/output"
LIBS_DIR="$PROJECT_DIR/libs"

# --- Auto-detect Android SDK build tools ---
find_build_tools() {
    local sdk_root="${ANDROID_SDK_ROOT:-$HOME/Library/Android/sdk}"
    local bt
    bt=$(find "$sdk_root/build-tools" -maxdepth 1 -type d 2>/dev/null | sort -V | tail -1)
    if [ -n "$bt" ] && [ -d "$bt" ]; then
        echo "$bt"; return
    fi
    echo ""
}

BUILD_TOOLS="${ANDROID_SDK_BUILD_TOOLS:-$(find_build_tools)}"
if [ -z "$BUILD_TOOLS" ] || [ ! -d "$BUILD_TOOLS" ]; then
    echo "❌ Android SDK build-tools not found."
    echo "   Set ANDROID_SDK_BUILD_TOOLS to your build-tools path."
    exit 1
fi

ZIPALIGN="$BUILD_TOOLS/zipalign"
APKSIGNER="$BUILD_TOOLS/apksigner"

KEYSTORE="$PROJECT_DIR/config/rmhook.keystore"
KEYSTORE_PASS="rmhook"
KEY_ALIAS="rmhook"

# --- Portable in-place sed (GNU sed vs BSD/macOS sed take -i differently) ---
sed_inplace() {
    if sed --version >/dev/null 2>&1; then
        # GNU sed
        sed -i "$@"
    else
        # BSD/macOS sed requires an explicit (possibly empty) backup suffix
        sed -i '' "$@"
    fi
}

usage() {
    echo "Usage: $0 [apk_dir]"
    echo ""
    echo "  apk_dir  - Directory containing the split APKs (default: ~/reMarkable/app)"
    echo ""
    echo "Output: patched APKs in $OUTPUT_DIR/"
    exit 1
}

APK_DIR="${1:-$HOME/reMarkable/app}"

echo "💉 Injecting RMHook into reMarkable APKs..."
echo "📦 APK source: $APK_DIR"
echo ""

if [ ! -f "$BUILD_DIR/librmhook.so" ]; then
    echo "❌ librmhook.so not found. Run ./scripts/build.sh first."
    exit 1
fi

# Clean and create output directory
rm -rf "$OUTPUT_DIR"
mkdir -p "$OUTPUT_DIR"
rm -rf "$BUILD_DIR/base" "$BUILD_DIR/split_arm64" "$BUILD_DIR/split_xxhdpi"
rm -f "$BUILD_DIR"/*.apk

# --- Step 1: Decompile and patch base.apk (smali injection only, no resources) ---
echo "[1/4] Patching base.apk (smali injection)..."
# Use -r (--no-res) to skip decoding/recoding resources. We only need to modify the DEX.
apktool d -r "$APK_DIR/base.apk" -o "$BUILD_DIR/base" -f

MAIN_SMALI="$BUILD_DIR/base/smali/com/remarkable/mobile/MainActivity.smali"
if [ ! -f "$MAIN_SMALI" ]; then
    echo "❌ MainActivity.smali not found"
    exit 1
fi

# Insert System.loadLibrary("rmhook") at the very start of <clinit>
sed_inplace '/.method static constructor <clinit>()V/{
n
/.locals/{
a\
\
    const-string v0, "rmhook"\
\
    invoke-static {v0}, Ljava/lang/System;->loadLibrary(Ljava/lang/String;)V
}
}' "$MAIN_SMALI"

if ! grep -q 'invoke-static {v0}, Ljava/lang/System;->loadLibrary' "$MAIN_SMALI"; then
    echo "❌ Failed to inject loadLibrary(\"rmhook\") into MainActivity.smali"
    echo "   The <clinit> block's shape may have changed - inspect $MAIN_SMALI manually."
    exit 1
fi

echo "✅ Injected System.loadLibrary(\"rmhook\") into MainActivity.<clinit>"

apktool b "$BUILD_DIR/base" -o "$BUILD_DIR/base_patched.apk"

# --- Step 2: Patch split_config.arm64_v8a.apk (add hook library) ---
echo "[2/4] Injecting ShadowHook libraries into arm64 split..."
SPLIT_ARM64="$APK_DIR/split_config.arm64_v8a.apk"
if [ ! -f "$SPLIT_ARM64" ]; then
    echo "❌ $SPLIT_ARM64 not found"
    exit 1
fi

mkdir -p "$BUILD_DIR/split_arm64"
unzip -q "$SPLIT_ARM64" -d "$BUILD_DIR/split_arm64"

cp "$BUILD_DIR/librmhook.so" "$BUILD_DIR/split_arm64/lib/arm64-v8a/"
cp "$LIBS_DIR/libshadowhook_nothing.so" "$BUILD_DIR/split_arm64/lib/arm64-v8a/"

rm -rf "$BUILD_DIR/split_arm64/META-INF"

# Store everything uncompressed (-0), including the native libraries.
(cd "$BUILD_DIR/split_arm64" && zip -q -r -0 "$BUILD_DIR/split_arm64_raw.apk" .)

# --- Step 3: Align APKs ---
echo "[3/4] Aligning APKs..."
"$ZIPALIGN" -p -f 4 "$BUILD_DIR/base_patched.apk" "$BUILD_DIR/base_aligned.apk"
"$ZIPALIGN" -p -f 4 "$BUILD_DIR/split_arm64_raw.apk" "$BUILD_DIR/split_arm64_aligned.apk"

SPLIT_XXHDPI="$APK_DIR/split_config.xxhdpi.apk"
if [ -f "$SPLIT_XXHDPI" ]; then
    mkdir -p "$BUILD_DIR/split_xxhdpi"
    unzip -q "$SPLIT_XXHDPI" -d "$BUILD_DIR/split_xxhdpi"
    rm -rf "$BUILD_DIR/split_xxhdpi/META-INF"
    (cd "$BUILD_DIR/split_xxhdpi" && zip -q -r -0 "$BUILD_DIR/split_xxhdpi_raw.apk" .)
    "$ZIPALIGN" -p -f 4 "$BUILD_DIR/split_xxhdpi_raw.apk" "$BUILD_DIR/split_xxhdpi_aligned.apk"
fi

# --- Step 4: Sign all APKs ---
echo "[4/4] Signing APKs..."

if [ ! -f "$KEYSTORE" ]; then
    mkdir -p "$(dirname "$KEYSTORE")"
    echo "🔑 Generating signing keystore..."
    keytool -genkeypair -v \
        -keystore "$KEYSTORE" \
        -alias "$KEY_ALIAS" \
        -keyalg RSA -keysize 2048 \
        -validity 10000 \
        -storepass "$KEYSTORE_PASS" \
        -keypass "$KEYSTORE_PASS" \
        -dname "CN=RMHook, OU=Dev, O=RMHook, L=Unknown, ST=Unknown, C=US" \
        2>/dev/null
fi

sign_apk() {
    local input="$1"
    local output="$2"
    "$APKSIGNER" sign \
        --ks "$KEYSTORE" \
        --ks-pass "pass:$KEYSTORE_PASS" \
        --ks-key-alias "$KEY_ALIAS" \
        --key-pass "pass:$KEYSTORE_PASS" \
        --out "$output" \
        "$input"
}

sign_apk "$BUILD_DIR/base_aligned.apk" "$OUTPUT_DIR/base.apk"
sign_apk "$BUILD_DIR/split_arm64_aligned.apk" "$OUTPUT_DIR/split_config.arm64_v8a.apk"
if [ -f "$BUILD_DIR/split_xxhdpi_aligned.apk" ]; then
    sign_apk "$BUILD_DIR/split_xxhdpi_aligned.apk" "$OUTPUT_DIR/split_config.xxhdpi.apk"
fi

echo ""
echo "✅ Injection successful!"
echo "📍 Patched APKs in: $OUTPUT_DIR/"
ls -lh "$OUTPUT_DIR/"*.apk
echo ""
echo "🚀 To install on your reMarkable:"
echo "   adb uninstall com.remarkable.mobile"
echo "   adb install-multiple --no-streaming $OUTPUT_DIR/base.apk $OUTPUT_DIR/split_config.arm64_v8a.apk $OUTPUT_DIR/split_config.xxhdpi.apk"
echo "   adb push rmhook.conf /sdcard/Android/data/com.remarkable.mobile/files/rmhook.conf"
echo ""
echo "🔍 To debug:"
echo "   adb logcat -s 'RMHook:*'"
echo ""