#!/bin/bash
#
# setup.sh — Persona_CoreIKDemo: one-time setup per reboot
# - Niall George Horn; June 2026

# Info: 
# What this does:
#   1. Verifies Xcode is installed and the xrsimulator SDK is available
#   2. Discovers all installed visionOS simulator runtimes on this Mac
#   3. Picks one (the newest by default; override with --version <num>)
#   4. Copies CoreIK.framework (binary + Info.plist + the 3 rig files the
#      3-point solver needs) to /tmp/CoreIK.framework/ so we load from a
#      space-free path. (CoreIK's CFURLCopyPath URL-encodes spaces → %20 →
#      std::ifstream::open fails silently. /tmp sidesteps this.)
#   5. Ad-hoc codesigns the /tmp copy so dyld accepts it
#   6. Records the build tag at /tmp/CoreIK.framework/BUILD so the
#      Swift app and CLI demo can look up the correct per-build
#      function offsets at runtime
#   7. Prints the Apple Vision Pro simulator UDID(s) for that version
#
# Usage:
#   ./deploy/setup.sh                    # default: newest installed xrOS
#   ./deploy/setup.sh --version 26.2     # pick a specific version
#   ./deploy/setup.sh --list             # just list what's installed and exit
#
# /tmp is cleared on reboot — re-run this script after every reboot.
# Safe to re-run — /tmp/CoreIK.framework is overwritten.
# Claude saved my arse here, not going to lie! 
#

set -euo pipefail

# =========== Colours ===========
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
CYAN='\033[0;36m'; BOLD='\033[1m';   RESET='\033[0m'

ok()   { echo -e "${GREEN}  ✅ $*${RESET}"; }
info() { echo -e "${CYAN}  ℹ  $*${RESET}"; }
warn() { echo -e "${YELLOW}  ⚠  $*${RESET}"; }
fail() { echo -e "${RED}  ❌ $*${RESET}"; exit 1; }

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" && pwd)"

echo ""
echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${RESET}"
echo -e "${BOLD}  Persona_CoreIKDemo — Setup${RESET}"
echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${RESET}"
echo ""

# =========== Parse args ===========
#
# Version selection priority (highest to lowest):
#   1. --version <N.N> flag
#   2. COREIK_VERSION env var  (e.g. export COREIK_VERSION=26.5 in ~/.zshrc like I did for Tribunal demo)
#   3. Auto: newest installed runtime
#
# The env var is useful for pinning a specific runtime during ABI verification
# across builds without having to pass --version every time.

WANT_VERSION="${COREIK_VERSION:-}"
LIST_ONLY=false

while [ $# -gt 0 ]; do
    case "$1" in
        --version) WANT_VERSION="$2"; shift 2 ;;
        --version=*) WANT_VERSION="${1#*=}"; shift ;;
        --list) LIST_ONLY=true; shift ;;
        -h|--help)
            grep -E '^#( |$)' "$0" | sed 's/^# \{0,1\}//'
            exit 0 ;;
        *) fail "Unknown argument: $1 (try --help)" ;;
    esac
done

[ -n "$WANT_VERSION" ] && info "Version pinned to: ${WANT_VERSION} (${COREIK_VERSION:+env var}${COREIK_VERSION:-flag})"

# =========== Step 1: Xcode present? ===========

echo -e "${BOLD}[1/5] Checking Xcode...${RESET}"

if ! command -v xcodebuild &>/dev/null; then
    fail "xcodebuild not found. Install Xcode from the App Store or Apple Developer portal."
fi

XCODE_VER=$(xcodebuild -version 2>/dev/null | head -1 || echo "unknown")
ok "xcodebuild found: $XCODE_VER"

if ! xcodebuild -showsdks 2>/dev/null | grep -q "xrsimulator"; then
    warn "xrsimulator SDK not found in active Xcode."
    warn "You may need: sudo xcode-select -s /Applications/Xcode-beta.app/Contents/Developer"
fi

echo ""

# =========== Step 2: Enumerate installed visionOS simulator runtimes ===========

echo -e "${BOLD}[2/5] Discovering installed visionOS simulator runtimes...${RESET}"

# We look at the mounted CoreSimulator volumes — one per installed xrOS runtime.
# Each volume contains a SystemVersion.plist we can mine for the build tag.
declare -a RUNTIMES   # entries: "OS_VERSION|BUILD_TAG|VOLUME_DIR|PRFW_PATH"

for vol_dir in /Library/Developer/CoreSimulator/Volumes/xrOS_*; do
    [ -d "$vol_dir" ] || continue

    # Find the xrOS N.N.simruntime profile dir inside this volume
    runtime_dir=$(find "$vol_dir/Library/Developer/CoreSimulator/Profiles/Runtimes" \
                       -maxdepth 1 -name "xrOS *.simruntime" -type d 2>/dev/null | head -1)
    [ -n "$runtime_dir" ] || continue

    sv_plist="$runtime_dir/Contents/Resources/RuntimeRoot/System/Library/CoreServices/SystemVersion.plist"
    [ -f "$sv_plist" ] || continue

    os_ver=$(plutil -extract ProductVersion raw -o - "$sv_plist" 2>/dev/null || echo "?")
    build=$( plutil -extract ProductBuildVersion raw -o - "$sv_plist" 2>/dev/null || echo "?")
    prfw="$runtime_dir/Contents/Resources/RuntimeRoot/System/Library/PrivateFrameworks"

    RUNTIMES+=("${os_ver}|${build}|${vol_dir}|${prfw}")
done

if [ ${#RUNTIMES[@]} -eq 0 ]; then
    # Volumes are mounted on demand by CoreSimulator. Opening Simulator.app
    # triggers the mount — wait up to 60 seconds for xrOS_* to appear.
    warn "No xrOS volumes mounted yet. Opening Simulator.app to trigger mount..."
    open -a Simulator 2>/dev/null || true
    WAITED=0
    until ls /Library/Developer/CoreSimulator/Volumes/xrOS_* &>/dev/null; do
        sleep 2; WAITED=$((WAITED + 2))
        if [ $WAITED -ge 60 ]; then
            fail "Timed out waiting for xrOS volumes to mount.
    Make sure at least one visionOS Simulator runtime is installed:
    Xcode > Settings > Platforms > visionOS"
        fi
    done
    info "Volumes mounted after ${WAITED}s. Re-scanning..."
    for vol_dir in /Library/Developer/CoreSimulator/Volumes/xrOS_*; do
        [ -d "$vol_dir" ] || continue
        runtime_dir=$(find "$vol_dir/Library/Developer/CoreSimulator/Profiles/Runtimes" \
                           -maxdepth 1 -name "xrOS *.simruntime" -type d 2>/dev/null | head -1)
        [ -n "$runtime_dir" ] || continue
        sv_plist="$runtime_dir/Contents/Resources/RuntimeRoot/System/Library/CoreServices/SystemVersion.plist"
        [ -f "$sv_plist" ] || continue
        os_ver=$(plutil -extract ProductVersion raw -o - "$sv_plist" 2>/dev/null || echo "?")
        build=$( plutil -extract ProductBuildVersion raw -o - "$sv_plist" 2>/dev/null || echo "?")
        prfw="$runtime_dir/Contents/Resources/RuntimeRoot/System/Library/PrivateFrameworks"
        RUNTIMES+=("${os_ver}|${build}|${vol_dir}|${prfw}")
    done
    if [ ${#RUNTIMES[@]} -eq 0 ]; then
        fail "No visionOS simulator runtimes found even after mount.
    Install one via Xcode > Settings > Platforms > visionOS (~6.7 GB)."
    fi
fi

# Pretty-print the list
echo ""
printf "  ${BOLD}%-10s %-10s %s${RESET}\n" "Version" "Build" "Volume"
printf "  %-10s %-10s %s\n" "-------" "-----" "------"
for entry in "${RUNTIMES[@]}"; do
    IFS='|' read -r os_ver build vol _ <<< "$entry"
    printf "  ${CYAN}%-10s %-10s${RESET} %s\n" "$os_ver" "$build" "$vol"
done
echo ""

if $LIST_ONLY; then exit 0; fi

# =========== Step 3: Pick which runtime to target ===========

echo -e "${BOLD}[3/5] Selecting target runtime...${RESET}"

PICKED=""
if [ -n "$WANT_VERSION" ]; then
    # Match against ProductVersion (e.g. "26.2"). Tolerate build-tag matches too.
    for entry in "${RUNTIMES[@]}"; do
        IFS='|' read -r os_ver build _ _ <<< "$entry"
        if [ "$os_ver" = "$WANT_VERSION" ] || [ "$build" = "$WANT_VERSION" ]; then
            PICKED="$entry"
            break
        fi
    done
    if [ -z "$PICKED" ]; then
        fail "No installed runtime matches --version=$WANT_VERSION"
    fi
else
    # Default: newest by ProductVersion. sort -V handles 26.2 < 26.10 correctly.
    PICKED=$(printf '%s\n' "${RUNTIMES[@]}" | sort -t'|' -k1,1 -V | tail -1)
fi

IFS='|' read -r OS_VERSION BUILD_TAG VOL_DIR PRFW <<< "$PICKED"
COREIK_SRC="${PRFW}/CoreIK.framework"

ok "Target: visionOS ${OS_VERSION} (build ${BUILD_TAG})"
info "CoreIK source: ${COREIK_SRC}"

if [ ! -f "${COREIK_SRC}/CoreIK" ]; then
    fail "CoreIK binary not found at ${COREIK_SRC}/CoreIK (runtime may be corrupted)"
fi

# Check whether we have pre-generated offsets for this build.
# Offsets for 26.2 (23N301), 26.4.1 (23O249a), and 26.5 (23O470) are committed
# in deploy/coreik_offsets/ — those builds work out of the box.
# Any newer build (visionOS 27.x etc) will need offsets derived first.
OFFSETS_FILE="${SCRIPT_DIR}/coreik_offsets/${BUILD_TAG}.txt"
if [ ! -f "$OFFSETS_FILE" ]; then
    echo ""
    echo -e "${RED}${BOLD}  ┌─────────────────────────────────────────────────────┐${RESET}"
    echo -e "${RED}${BOLD}  │  NEW BUILD — offsets required before app will work  │${RESET}"
    echo -e "${RED}${BOLD}  └─────────────────────────────────────────────────────┘${RESET}"
    warn "No pre-generated offsets for build ${BUILD_TAG}."
    warn "The app will launch but the solver will not initialise."
    warn ""
    warn "Run this now (takes ~5 seconds):"
    warn "    ${SCRIPT_DIR}/derive_offsets.sh /tmp/CoreIK.framework/CoreIK"
    warn ""
    warn "Then paste the printed Swift block into:"
    warn "    HelloVisionOS/CoreIKBridge.swift  (inside kCoreIKOffsets)"
    warn "And commit the generated file:"
    warn "    deploy/coreik_offsets/${BUILD_TAG}.txt"
    echo ""
else
    ok "Known offsets available for build ${BUILD_TAG} ✓"
fi

echo ""

# =========== Step 4: Populate /tmp/CoreIK.framework ===========

echo -e "${BOLD}[4/5] Setting up /tmp/CoreIK.framework/ ...${RESET}"
info "Copying from: ${COREIK_SRC}"
info "Copying to  : /tmp/CoreIK.framework"

# Clean-slate any stale copy from a previous run
rm -rf /tmp/CoreIK.framework
mkdir -p /tmp/CoreIK.framework

cp "${COREIK_SRC}/CoreIK" /tmp/CoreIK.framework/CoreIK
ok "Copied CoreIK binary"

if [ -f "${COREIK_SRC}/Info.plist" ]; then
    cp "${COREIK_SRC}/Info.plist" /tmp/CoreIK.framework/Info.plist
    ok "Copied Info.plist"
fi

# Only the three rig files the 3-point solver actually needs.
# The runtime also ships these rigs (not copied here, but available in COREIK_SRC
# if you want to experiment with more input joints / end-effectors in future):
#
#   BodyFusion_FullBodyOrion.ikn      — full-body Orion solver (legs + torso)
#   BodyFusion_FullBodyOrionV2.ikn    — Orion V2, likely updated propagation filter
#   BodyFusion_FullBodyArvinoD.ikn    — Arvino solver variant (Arvino = Apple's IK brand)
#   BodyFusion_FullBodyCombi.ikn      — combined Arvino + Orion rig
#   BodyFusion_LeftHand.ikn           — hand tracking rig, left
#   BodyFusion_RightHand.ikn          — hand tracking rig, right
#   BodyFusion_ArmPoses.ikn           — arm pose priors (used in full-body mode)
#   BodyFusion_HandLockPoses.ikn      — hand-lock constraint poses
#
# To test full-body reconstruction: copy the FullBody* rigs here too, then wire
# a different CreateSolver call with the appropriate rig type index.
NEEDED_IKNS=(
    "BodyFusion_PartialBody3PointRot.ikn"   # active rig: head + 2 wrists → 91 joints
    "BodyFusion_RestPose.ikn"               # default T-pose, loaded at solver init
    "EnrollmentRigV2.ikn"                   # enrollment geometry, required by CreateSolver
)

IKN_COUNT=0
for ikn_name in "${NEEDED_IKNS[@]}"; do
    src="${COREIK_SRC}/${ikn_name}"
    if [ -f "$src" ]; then
        cp "$src" /tmp/CoreIK.framework/
        IKN_COUNT=$((IKN_COUNT + 1))
    else
        fail "Required rig file missing from runtime: ${ikn_name}"
    fi
done

ok "Copied ${IKN_COUNT} .ikn rig files (PartialBody3PointRot + RestPose + EnrollmentRigV2)"

# Write BUILD metadata — CoreIKBridge.swift + main.c read this to pick offsets
printf '%s\n' "$BUILD_TAG" > /tmp/CoreIK.framework/BUILD
printf 'BUILD=%s\nOS_VERSION=%s\nSOURCE=%s\nSETUP_AT=%s\n' \
    "$BUILD_TAG" "$OS_VERSION" "$COREIK_SRC" "$(date -u +%Y-%m-%dT%H:%M:%SZ)" \
    > /tmp/CoreIK.framework/SETUP.info
ok "Wrote /tmp/CoreIK.framework/BUILD = ${BUILD_TAG}"

echo ""

# =========== Step 5: Ad-hoc codesign ===========

echo -e "${BOLD}[5/5] Codesigning /tmp/CoreIK.framework/CoreIK ...${RESET}"

if command -v codesign &>/dev/null; then
    codesign --force --sign - /tmp/CoreIK.framework/CoreIK 2>/dev/null
    ok "Codesigned (ad-hoc)"
else
    warn "codesign not found — dyld may reject the copy"
fi

echo ""

# =========== Summary ===========

echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${RESET}"
echo -e "${BOLD}  /tmp/CoreIK.framework contents:${RESET}"
ls -la /tmp/CoreIK.framework/ | tail -n +2 | while IFS= read -r line; do
    echo "    $line"
done

echo ""
echo -e "${BOLD}  Apple Vision Pro Simulator UDIDs (${OS_VERSION}):${RESET}"
xcrun simctl list devices 2>/dev/null \
    | awk -v ver="com.apple.CoreSimulator.SimRuntime.xrOS-${OS_VERSION//./-}" \
          '/^-- /{rt=$0} $0 ~ "Apple Vision Pro" && rt ~ ver {print}' \
    | head -5 \
    | while IFS= read -r line; do echo -e "    ${CYAN}${line}${RESET}"; done

# Fall back: print all Vision Pro UDIDs if the filter above found nothing
# (the xrOS runtime header format changes across Xcode versions)
if ! xcrun simctl list devices 2>/dev/null | grep -q "Apple Vision Pro.*Booted\|Apple Vision Pro.*Shutdown"; then
    :
else
    xcrun simctl list devices 2>/dev/null | grep "Apple Vision Pro" | head -5 | while IFS= read -r line; do
        echo -e "    ${CYAN}${line}${RESET}" 2>/dev/null || true
    done
fi

echo ""
echo -e "${GREEN}${BOLD}  ✅ Setup complete.${RESET}"
echo ""
echo -e "  Next steps:"
echo -e "    ${CYAN}./deploy/build.sh${RESET}     ← compile the app for ${OS_VERSION}"
echo -e "    ${CYAN}./deploy/launch.sh${RESET}    ← boot simulator, install, launch"
echo -e "    ${CYAN}./deploy/logs.sh${RESET}      ← stream live CoreIK invocation log"
echo ""
echo -e "  Or run the full end-to-end check:"
echo -e "    ${CYAN}./deploy/verify.sh${RESET}"
echo ""
echo -e "${BOLD}${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${RESET}"
echo ""
