#!/bin/bash

# =============================================================
# Repository + Submodules Update Script
# =============================================================
# - Run this script from the root of the main repository.
# - Submodules and their branches are read from .gitmodules.
# - For each submodule, a normal pull is attempted first.
# - If the pull fails, the user is asked whether to force-reset.
# =============================================================

# Branch to pull on the main repository
MAIN_BRANCH="main"

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

# Counters
SUCCESS_COUNT=0
RESET_COUNT=0
FAIL_COUNT=0
SKIPPED_COUNT=0
FAILED_ITEMS=()

# =============================================================
# Determine main repository root automatically
# =============================================================
MAIN_REPO=$(git rev-parse --show-toplevel 2>/dev/null)
if [ -z "$MAIN_REPO" ]; then
    echo -e "${RED}ERROR: This script must be run inside a Git repository.${NC}"
    exit 1
fi

GITMODULES_FILE="${MAIN_REPO}/.gitmodules"
if [ ! -f "$GITMODULES_FILE" ]; then
    echo -e "${RED}ERROR: No .gitmodules file found at ${GITMODULES_FILE}${NC}"
    exit 1
fi

# =============================================================
# Update the main repository (normal pull)
# =============================================================
update_main_repo() {
    echo ""
    echo -e "${CYAN}=============================================================${NC}"
    echo -e "${CYAN}Main repository: ${MAIN_REPO}${NC}"
    echo -e "${CYAN}Branch:          ${MAIN_BRANCH}${NC}"
    echo -e "${CYAN}=============================================================${NC}"

    cd "$MAIN_REPO" || {
        echo -e "${RED}ERROR: Cannot enter directory: ${MAIN_REPO}${NC}"
        return 1
    }

    # Stash uncommitted changes (excluding submodule pointer changes) so the pull can proceed
    local stashed=0
    if ! git diff --quiet --ignore-submodules || ! git diff --cached --quiet --ignore-submodules; then
        echo -e "${YELLOW}WARNING: Uncommitted changes detected in main repo. Stashing temporarily...${NC}"
        if git stash push --include-untracked -m "update-submodules.sh: auto-stash" 2>/dev/null; then
            stashed=1
        else
            echo -e "${RED}ERROR: Could not stash uncommitted changes. Skipping main repo pull.${NC}"
            return 1
        fi
    fi

    echo -e "${YELLOW}Fetching from remote...${NC}"
    if ! git fetch --prune; then
        echo -e "${RED}ERROR: git fetch failed${NC}"
        return 1
    fi

    echo -e "${YELLOW}Checking out branch '${MAIN_BRANCH}'...${NC}"
    if ! git checkout "$MAIN_BRANCH"; then
        echo -e "${RED}ERROR: Could not checkout branch '${MAIN_BRANCH}'${NC}"
        return 1
    fi

    echo -e "${YELLOW}Pulling latest changes (fast-forward only)...${NC}"
    if ! git pull --ff-only origin "$MAIN_BRANCH"; then
        echo -e "${RED}ERROR: git pull failed${NC}"
        return 1
    fi

    # Sync submodule URLs and initialize any new ones (without updating contents)
    echo -e "${YELLOW}Synchronizing submodule URLs...${NC}"
    git submodule sync --recursive

    echo -e "${YELLOW}Initializing submodules (if any new)...${NC}"
    git submodule init

    if [ $stashed -eq 1 ]; then
        echo -e "${YELLOW}Restoring stashed changes...${NC}"
        if git stash pop; then
            echo -e "${GREEN}Stashed changes restored.${NC}"
        else
            echo -e "${YELLOW}WARNING: Could not auto-restore stashed changes (possible conflicts).${NC}"
            echo -e "${YELLOW}Run 'git stash pop' manually to restore them.${NC}"
        fi
    fi

    echo -e "${GREEN}SUCCESS: Main repository updated.${NC}"
    return 0
}

# =============================================================
# Force-reset a single submodule to the latest remote branch
# =============================================================
force_reset_submodule() {
    local sub_path="$1"
    local branch="$2"
    local full_path="${MAIN_REPO}/${sub_path}"

    cd "$full_path" || return 1

    # Cannot reset to a branch that doesn't exist on the remote
    if ! git rev-parse --verify "origin/${branch}" &>/dev/null; then
        echo -e "${YELLOW}Remote branch 'origin/${branch}' not found; skipping force reset.${NC}"
        return 2
    fi

    # Abort any in-progress merge/rebase/cherry-pick and clear a conflicted index
    git merge --abort 2>/dev/null || true
    git cherry-pick --abort 2>/dev/null || true
    git rebase --abort 2>/dev/null || true

    echo -e "${YELLOW}Resetting to origin/${branch}...${NC}"
    if ! git reset --hard "origin/${branch}"; then
        echo -e "${RED}ERROR: git reset failed${NC}"
        return 1
    fi

    echo -e "${YELLOW}Cleaning untracked files...${NC}"
    git clean -fdx

    return 0
}

# =============================================================
# Update a single submodule (try normal pull first)
# =============================================================
update_submodule() {
    local sub_path="$1"
    local branch="$2"
    local full_path="${MAIN_REPO}/${sub_path}"

    echo ""
    echo -e "${CYAN}-------------------------------------------------------------${NC}"
    echo -e "${CYAN}Submodule: ${sub_path}${NC}"
    echo -e "${CYAN}Branch:    ${branch}${NC}"
    echo -e "${CYAN}-------------------------------------------------------------${NC}"

    if [ ! -d "$full_path" ]; then
        echo -e "${YELLOW}Submodule '${sub_path}' directory missing. Attempting to initialize...${NC}"
        cd "$MAIN_REPO" || return 1
        if git submodule update --init -- "$sub_path" 2>/dev/null; then
            echo -e "${GREEN}SUCCESS: Submodule '${sub_path}' initialized.${NC}"
            SUCCESS_COUNT=$((SUCCESS_COUNT + 1))
            return 0
        else
            echo -e "${YELLOW}WARNING: Could not initialize '${sub_path}' — it may not be committed to the git index.${NC}"
            echo -e "${YELLOW}         Run 'git submodule add <url> ${sub_path}' and commit to register it properly.${NC}"
            SKIPPED_COUNT=$((SKIPPED_COUNT + 1))
            FAILED_ITEMS+=("${sub_path} (not in git index)")
            return 0
        fi
    fi

    cd "$full_path" || {
        echo -e "${RED}ERROR: Cannot enter directory: ${full_path}${NC}"
        FAIL_COUNT=$((FAIL_COUNT + 1))
        FAILED_ITEMS+=("${sub_path} (cd failed)")
        return 1
    }

    # Stash any uncommitted changes so the pull can proceed cleanly
    local stashed=0
    if ! git diff --quiet || ! git diff --cached --quiet; then
        echo -e "${YELLOW}Uncommitted changes detected. Stashing temporarily...${NC}"
        if git stash push --include-untracked -m "update-submodules.sh: auto-stash" 2>/dev/null; then
            stashed=1
        else
            echo -e "${YELLOW}WARNING: Could not stash changes; proceeding anyway...${NC}"
        fi
    fi

    # Prune stale/corrupt remote-tracking refs before fetching
    git remote prune origin 2>/dev/null || true

    # Fetch latest refs (retry once after pruning if first attempt fails)
    echo -e "${YELLOW}Fetching from remote...${NC}"
    if ! git fetch --prune origin; then
        echo -e "${YELLOW}Fetch failed; retrying after pruning stale refs...${NC}"
        git remote prune origin 2>/dev/null || true
        if ! git fetch --prune origin; then
            echo -e "${RED}ERROR: git fetch failed${NC}"
            FAIL_COUNT=$((FAIL_COUNT + 1))
            FAILED_ITEMS+=("${sub_path} (fetch failed)")
            return 1
        fi
    fi

    # Verify the remote branch actually exists before attempting checkout/pull
    if ! git rev-parse --verify "origin/${branch}" &>/dev/null; then
        echo -e "${YELLOW}WARNING: Remote branch 'origin/${branch}' not found in '${sub_path}'. Skipping.${NC}"
        echo -e "${YELLOW}Check the 'branch' value in .gitmodules for this submodule.${NC}"
        SKIPPED_COUNT=$((SKIPPED_COUNT + 1))
        FAILED_ITEMS+=("${sub_path} (remote branch 'origin/${branch}' not found)")
        return 0
    fi

    # Try a normal (non-destructive) update
    local pull_failed=0

    echo -e "${YELLOW}Attempting normal checkout of '${branch}'...${NC}"
    if ! git checkout "$branch" 2>/dev/null; then
        echo -e "${YELLOW}Local branch '${branch}' not present or checkout failed.${NC}"
        pull_failed=1
    else
        echo -e "${YELLOW}Attempting fast-forward pull...${NC}"
        if ! git pull --ff-only origin "$branch"; then
            echo -e "${YELLOW}Fast-forward pull failed.${NC}"
            pull_failed=1
        fi
    fi

    if [ $pull_failed -eq 0 ]; then
        if [ $stashed -eq 1 ]; then
            echo -e "${YELLOW}Restoring stashed changes...${NC}"
            if ! git stash pop; then
                echo -e "${YELLOW}WARNING: Could not restore stashed changes in '${sub_path}'.${NC}"
                echo -e "${YELLOW}Run 'git stash pop' inside '${sub_path}' manually.${NC}"
            fi
        fi
        echo -e "${GREEN}SUCCESS: Submodule '${sub_path}' updated normally.${NC}"
        SUCCESS_COUNT=$((SUCCESS_COUNT + 1))
        return 0
    fi

    # Normal pull failed: ask the user about a force reset
    echo ""
    echo -e "${YELLOW}Normal update failed for submodule '${sub_path}'.${NC}"
    echo -e "${YELLOW}A force reset will:${NC}"
    echo -e "${YELLOW}  - Hard-reset the submodule to origin/${branch}${NC}"
    echo -e "${YELLOW}  - Delete any local changes${NC}"
    echo -e "${YELLOW}  - Remove untracked and ignored files${NC}"

    local answer=""
    # Read from the terminal explicitly so it works even when stdin is redirected
    if [ -t 0 ]; then
        read -r -p "Do you want to force-reset this submodule? [y/N] " answer
    else
        read -r -p "Do you want to force-reset this submodule? [y/N] " answer < /dev/tty
    fi

    case "$answer" in
        [Yy]|[Yy][Ee][Ss])
            force_reset_submodule "$sub_path" "$branch"
            local reset_rc=$?
            if [ $reset_rc -eq 0 ]; then
                [ $stashed -eq 1 ] && git stash drop 2>/dev/null
                echo -e "${GREEN}SUCCESS: Submodule '${sub_path}' force-reset to origin/${branch}.${NC}"
                RESET_COUNT=$((RESET_COUNT + 1))
                return 0
            elif [ $reset_rc -eq 2 ]; then
                echo -e "${YELLOW}Skipped '${sub_path}': remote branch 'origin/${branch}' not available.${NC}"
                SKIPPED_COUNT=$((SKIPPED_COUNT + 1))
                FAILED_ITEMS+=("${sub_path} (remote branch 'origin/${branch}' not found)")
                return 0
            else
                echo -e "${RED}ERROR: Force reset failed for '${sub_path}'.${NC}"
                FAIL_COUNT=$((FAIL_COUNT + 1))
                FAILED_ITEMS+=("${sub_path} (reset failed)")
                return 1
            fi
            ;;
        *)
            echo -e "${YELLOW}Skipped force reset for '${sub_path}'.${NC}"
            SKIPPED_COUNT=$((SKIPPED_COUNT + 1))
            FAILED_ITEMS+=("${sub_path} (skipped by user)")
            return 1
            ;;
    esac
}

# =============================================================
# Stage and commit all pending changes in the main repository
# =============================================================
commit_main_repo() {
    cd "$MAIN_REPO" || return 1

    if git diff --quiet && git diff --cached --quiet; then
        echo -e "${YELLOW}No changes to commit in main repository.${NC}"
        return 0
    fi

    echo -e "${YELLOW}Staging all changes in main repository...${NC}"
    git add -A

    echo -e "${YELLOW}Committing...${NC}"
    if git commit -m "chore: update submodule references"; then
        echo -e "${GREEN}SUCCESS: Changes committed in main repository.${NC}"
        return 0
    else
        echo -e "${RED}ERROR: git commit failed in main repository.${NC}"
        return 1
    fi
}

# =============================================================
# Read submodule list from .gitmodules
# =============================================================
read_submodules_from_gitmodules() {
    # Outputs lines: "<path>|<branch>"
    local names
    names=$(git config -f "$GITMODULES_FILE" --name-only --get-regexp '^submodule\..*\.path$' \
            | sed -e 's/^submodule\.//' -e 's/\.path$//')

    if [ -z "$names" ]; then
        return 0
    fi

    while IFS= read -r name; do
        [ -z "$name" ] && continue
        local sub_path
        local branch
        sub_path=$(git config -f "$GITMODULES_FILE" --get "submodule.${name}.path")
        branch=$(git config -f "$GITMODULES_FILE" --get "submodule.${name}.branch" 2>/dev/null)

        # Default branch if not specified in .gitmodules
        if [ -z "$branch" ]; then
            branch="main"
            echo -e "${YELLOW}NOTE: No 'branch' set for submodule '${name}'; defaulting to '${branch}'.${NC}" >&2
        fi

        echo "${sub_path}|${branch}"
    done <<< "$names"
}

# =============================================================
# Main flow
# =============================================================
echo -e "${CYAN}Starting repository + submodules update...${NC}"
echo -e "${CYAN}Main repository detected at: ${MAIN_REPO}${NC}"

START_DIR=$(pwd)

# Step 1: Update main repository
if ! update_main_repo; then
    echo -e "${RED}Main repository update failed. Aborting.${NC}"
    cd "$START_DIR" || exit 1
    exit 1
fi

# Step 2: Read submodules from .gitmodules
mapfile -t SUBMODULES < <(read_submodules_from_gitmodules)

if [ ${#SUBMODULES[@]} -eq 0 ]; then
    echo ""
    echo -e "${YELLOW}No submodules found in .gitmodules. Nothing more to do.${NC}"
    cd "$START_DIR" || exit 1
    exit 0
fi

echo ""
echo -e "${CYAN}=============================================================${NC}"
echo -e "${CYAN}Updating ${#SUBMODULES[@]} submodule(s) from .gitmodules...${NC}"
echo -e "${CYAN}=============================================================${NC}"

# Step 3: Update each submodule
for entry in "${SUBMODULES[@]}"; do
    sub_path="${entry%%|*}"
    branch="${entry##*|}"

    update_submodule "$sub_path" "$branch"

    cd "$START_DIR" || exit 1
done

# Step 4: Commit all resulting changes in the main repository
echo ""
echo -e "${CYAN}=============================================================${NC}"
echo -e "${CYAN}Committing changes in main repository...${NC}"
echo -e "${CYAN}=============================================================${NC}"
commit_main_repo

# =============================================================
# Summary
# =============================================================
echo ""
echo -e "${CYAN}=============================================================${NC}"
echo -e "${CYAN}                          SUMMARY${NC}"
echo -e "${CYAN}=============================================================${NC}"
echo -e "${GREEN}Updated normally:    ${SUCCESS_COUNT}${NC}"
echo -e "${GREEN}Force-reset:         ${RESET_COUNT}${NC}"
echo -e "${YELLOW}Skipped by user:     ${SKIPPED_COUNT}${NC}"
echo -e "${RED}Failed:              ${FAIL_COUNT}${NC}"

if [ ${#FAILED_ITEMS[@]} -gt 0 ]; then
    echo ""
    echo -e "${YELLOW}Submodules not updated normally:${NC}"
    for failed in "${FAILED_ITEMS[@]}"; do
        echo -e "${YELLOW}  - ${failed}${NC}"
    done
fi

if [ ${FAIL_COUNT} -gt 0 ]; then
    exit 1
fi

echo ""
echo -e "${GREEN}Done.${NC}"
exit 0