#!/bin/bash
# Apply Local Repository State
# Applies captured git state including local branches, staged/unstaged changes, untracked files

set -o pipefail

STATE_FILE="${LOCAL_STATE_FILE:-/task-output/local-state.json}"
WORKSPACE_DIR="${WORKSPACE_DIR:-/workspace}"

# Track warnings and errors
HAS_WARNINGS=false
WARNING_MESSAGES=()

# Logging functions
log() {
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] $*" >&2
}

log_error() {
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] ERROR: $*" >&2
    HAS_WARNINGS=true
    WARNING_MESSAGES+=("ERROR: $*")
}

log_warning() {
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] WARNING: $*" >&2
    HAS_WARNINGS=true
    WARNING_MESSAGES+=("WARNING: $*")
}

log_success() {
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] SUCCESS: $*" >&2
}

# Check if state file exists
if [ ! -f "$STATE_FILE" ]; then
    log "No local state file found at $STATE_FILE, skipping state application"
    exit 0
fi

log "════════════════════════════════════════════════════════════════"
log "  Applying Local Repository State"
log "════════════════════════════════════════════════════════════════"
log ""

# Validate JSON
if ! jq empty "$STATE_FILE" 2>/dev/null; then
    log_error "Invalid JSON in state file: $STATE_FILE"
    exit 1
fi

# Show capture info
CAPTURED_AT=$(jq -r '.captured_at' "$STATE_FILE")
REPOS_FOUND=$(jq -r '.repos_found | join(", ")' "$STATE_FILE")
REPOS_MISSING=$(jq -r '.repos_missing | join(", ")' "$STATE_FILE")

log "📋 State Information:"
log "   Captured at: $CAPTURED_AT"
log "   Repositories: $REPOS_FOUND"
if [ -n "$REPOS_MISSING" ] && [ "$REPOS_MISSING" != "" ]; then
    log "   Missing (will use defaults): $REPOS_MISSING"
fi
log ""

# Process each repository
jq -r '.repositories | keys[]' "$STATE_FILE" | while read -r repo_name; do
    # Look up the path for this repo from REPOS_CONFIG
    repo_path_from_json=""
    if [ -n "$REPOS_CONFIG" ]; then
        repo_path_from_json=$(echo "$REPOS_CONFIG" | jq -r --arg name "$repo_name" '.[] | select(.name == $name) | .path')
    fi

    # Use path field from JSON if available, otherwise default to repo_name
    if [ -n "$repo_path_from_json" ] && [ "$repo_path_from_json" != "null" ] && [ "$repo_path_from_json" != "" ]; then
        REPO_PATH="$WORKSPACE_DIR/$repo_path_from_json"
    else
        REPO_PATH="$WORKSPACE_DIR/$repo_name"
    fi

    log "────────────────────────────────────────────────────────────────"
    log "📁 Repository: $repo_name"
    log "   Path: $REPO_PATH"
    log "────────────────────────────────────────────────────────────────"

    # Check if repo exists in workspace
    if [ ! -d "$REPO_PATH" ]; then
        log_error "Repository not found in workspace: $REPO_PATH"
        log "   Skipping..."
        log ""
        continue
    fi

    if [ ! -d "$REPO_PATH/.git" ]; then
        log_error "Not a git repository: $REPO_PATH"
        log "   Skipping..."
        log ""
        continue
    fi

    cd "$REPO_PATH" || {
        log_error "Failed to cd into $REPO_PATH"
        continue
    }

    # Extract repo state from JSON
    CURRENT_BRANCH=$(jq -r ".repositories.\"$repo_name\".current_branch" "$STATE_FILE")
    IS_REMOTE_TRACKING=$(jq -r ".repositories.\"$repo_name\".is_remote_tracking" "$STATE_FILE")
    BASE_BRANCH=$(jq -r ".repositories.\"$repo_name\".base_branch" "$STATE_FILE")
    COMMITS_AHEAD=$(jq -r ".repositories.\"$repo_name\".commits_ahead" "$STATE_FILE")

    log "   Current branch in container: $(git rev-parse --abbrev-ref HEAD)"
    log "   Target branch: $CURRENT_BRANCH"
    log "   Remote tracking: $IS_REMOTE_TRACKING"

    # Step 1: Handle branch setup
    if [ "$IS_REMOTE_TRACKING" = "true" ]; then
        # Branch exists on remote, checkout first
        log "   Checking out remote-tracking branch: $CURRENT_BRANCH"

        if git rev-parse --verify "origin/$CURRENT_BRANCH" >/dev/null 2>&1; then
            git checkout "$CURRENT_BRANCH" 2>&1 | sed 's/^/   /' || {
                log_error "Failed to checkout branch $CURRENT_BRANCH"
                continue
            }
        else
            log_error "Remote branch origin/$CURRENT_BRANCH not found"
            log "   Attempting to checkout local branch..."
            git checkout "$CURRENT_BRANCH" 2>&1 | sed 's/^/   /' || {
                log_error "Failed to checkout branch $CURRENT_BRANCH"
                continue
            }
        fi

        # Check for unpushed commits on remote-tracking branch
        if [ -n "$COMMITS_AHEAD" ] && [ "$COMMITS_AHEAD" != "null" ] && [ "$COMMITS_AHEAD" -gt 0 ]; then
            log "   📤 Unpushed commits detected: $COMMITS_AHEAD commit(s)"
            COMMIT_PATCH=$(jq -r ".repositories.\"$repo_name\".commit_patch" "$STATE_FILE")

            if [ -n "$COMMIT_PATCH" ] && [ "$COMMIT_PATCH" != "null" ]; then
                log "   Applying unpushed commits..."

                # Save patch to temp file
                TEMP_PATCH=$(mktemp)
                echo "$COMMIT_PATCH" > "$TEMP_PATCH"

                # Apply patches using git am
                if git am --3way < "$TEMP_PATCH" 2>&1 | sed 's/^/   /'; then
                    log_success "Successfully applied unpushed commits"
                else
                    log_error "Failed to apply commit patches, trying format-patch method..."
                    git am --abort 2>/dev/null || true

                    # Fallback: try applying as a single diff
                    if git apply < "$TEMP_PATCH" 2>&1 | sed 's/^/   /'; then
                        log_success "Applied changes as diff (commits not preserved)"
                    else
                        log_error "Failed to apply commit patches"
                    fi
                fi

                rm -f "$TEMP_PATCH"
            fi
        fi
    else
        # Local-only branch, need to recreate it
        log "   ⚠️  Local-only branch detected: $CURRENT_BRANCH"

        if [ -n "$BASE_BRANCH" ] && [ "$BASE_BRANCH" != "null" ]; then
            log "   Base branch: $BASE_BRANCH"
            log "   Commits ahead: $COMMITS_AHEAD"

            # Checkout base branch first
            if git rev-parse --verify "$BASE_BRANCH" >/dev/null 2>&1; then
                log "   Checking out base branch: $BASE_BRANCH"
                git checkout "$BASE_BRANCH" 2>&1 | sed 's/^/   /' || {
                    log_error "Failed to checkout base branch $BASE_BRANCH"
                    continue
                }
            else
                log_error "Base branch $BASE_BRANCH not found"
                continue
            fi

            # Create new branch
            log "   Creating local branch: $CURRENT_BRANCH"
            git checkout -b "$CURRENT_BRANCH" 2>&1 | sed 's/^/   /' || {
                # Branch might already exist
                git checkout "$CURRENT_BRANCH" 2>&1 | sed 's/^/   /' || {
                    log_error "Failed to create/checkout branch $CURRENT_BRANCH"
                    continue
                }
            }

            # Apply commit patches if available
            if [ "$COMMITS_AHEAD" != "null" ] && [ "$COMMITS_AHEAD" -gt 0 ]; then
                COMMIT_PATCH=$(jq -r ".repositories.\"$repo_name\".commit_patch" "$STATE_FILE")

                if [ -n "$COMMIT_PATCH" ] && [ "$COMMIT_PATCH" != "null" ]; then
                    log "   Applying $COMMITS_AHEAD commit(s)..."

                    # Save patch to temp file
                    TEMP_PATCH=$(mktemp)
                    echo "$COMMIT_PATCH" > "$TEMP_PATCH"

                    # Apply patches using git am
                    if git am --3way < "$TEMP_PATCH" 2>&1 | sed 's/^/   /'; then
                        log_success "Successfully applied commit patches"
                    else
                        log_error "Failed to apply commit patches, trying format-patch method..."
                        git am --abort 2>/dev/null || true

                        # Fallback: try applying as a single diff
                        if git apply < "$TEMP_PATCH" 2>&1 | sed 's/^/   /'; then
                            log_success "Applied changes as diff (commits not preserved)"
                        else
                            log_error "Failed to apply commit patches"
                        fi
                    fi

                    rm -f "$TEMP_PATCH"
                fi
            fi
        else
            log "   No base branch info, checking out branch as-is..."
            git checkout "$CURRENT_BRANCH" 2>&1 | sed 's/^/   /' || {
                log_error "Failed to checkout branch $CURRENT_BRANCH"
                continue
            }
        fi
    fi

    log_success "Branch setup complete"

    # Step 2: Apply unstaged changes (working directory changes)
    UNSTAGED=$(jq -r ".repositories.\"$repo_name\".unstaged" "$STATE_FILE")
    if [ -n "$UNSTAGED" ] && [ "$UNSTAGED" != "null" ] && [ "$UNSTAGED" != "" ]; then
        UNSTAGED_FILE_COUNT=$(echo "$UNSTAGED" | grep -c '^diff --git' || echo "0")
        log "   Applying unstaged changes ($UNSTAGED_FILE_COUNT files)..."

        TEMP_DIFF=$(mktemp)
        echo "$UNSTAGED" > "$TEMP_DIFF"

        if git apply --verbose < "$TEMP_DIFF" 2>&1 | sed 's/^/   /'; then
            log_success "Unstaged changes applied"
        else
            log_error "Failed to apply unstaged changes for $repo_name"
            log_warning "$repo_name: Unstaged changes had conflicts - check git status in container"
            log "   Attempting with --reject to show conflicts..."
            git apply --reject < "$TEMP_DIFF" 2>&1 | sed 's/^/   /' || true
        fi

        rm -f "$TEMP_DIFF"
    fi

    # Step 3: Apply staged changes (index)
    STAGED=$(jq -r ".repositories.\"$repo_name\".staged" "$STATE_FILE")
    if [ -n "$STAGED" ] && [ "$STAGED" != "null" ] && [ "$STAGED" != "" ]; then
        STAGED_FILE_COUNT=$(echo "$STAGED" | grep -c '^diff --git' || echo "0")
        log "   Applying staged changes ($STAGED_FILE_COUNT files)..."

        TEMP_STAGED=$(mktemp)
        echo "$STAGED" > "$TEMP_STAGED"

        if git apply --cached --verbose < "$TEMP_STAGED" 2>&1 | sed 's/^/   /'; then
            log_success "Staged changes applied"
        else
            log_error "Failed to apply staged changes for $repo_name"
            log_warning "$repo_name: Staged changes had conflicts - check git status in container"
            log "   Attempting with --reject to show conflicts..."
            git apply --cached --reject < "$TEMP_STAGED" 2>&1 | sed 's/^/   /' || true
        fi

        rm -f "$TEMP_STAGED"
    fi

    # Step 4: Create untracked files
    UNTRACKED_FILES=$(jq -r ".repositories.\"$repo_name\".untracked_files | keys[]" "$STATE_FILE" 2>/dev/null)
    if [ -n "$UNTRACKED_FILES" ]; then
        UNTRACKED_COUNT=$(echo "$UNTRACKED_FILES" | wc -l)
        log "   Creating $UNTRACKED_COUNT untracked file(s)..."

        echo "$UNTRACKED_FILES" | while read -r file_path; do
            if [ -z "$file_path" ]; then
                continue
            fi

            # Get file content from JSON (properly escaped)
            FILE_CONTENT=$(jq -r ".repositories.\"$repo_name\".untracked_files.\"$file_path\"" "$STATE_FILE")

            if [ -n "$FILE_CONTENT" ] && [ "$FILE_CONTENT" != "null" ]; then
                # Create parent directories if needed
                mkdir -p "$(dirname "$file_path")"

                # Write content to file
                echo "$FILE_CONTENT" > "$file_path"
                log "      + $file_path"
            fi
        done

        log_success "Untracked files created"
    fi

    # Step 5: Show final git status
    log ""
    log "   📊 Final repository state:"
    git status --short 2>&1 | sed 's/^/      /' || true
    log "   Current branch: $(git rev-parse --abbrev-ref HEAD)"
    log ""
    log_success "State application complete for $repo_name"
    log ""
done

log "════════════════════════════════════════════════════════════════"
log_success "Local state application complete!"
log "════════════════════════════════════════════════════════════════"
log ""

# Write summary for CLI to check
SUMMARY_FILE="/task-output/.local-state-summary"

if [ "$HAS_WARNINGS" = "true" ]; then
  # Build warnings JSON array
  WARNINGS_JSON="["
  first=true
  for msg in "${WARNING_MESSAGES[@]}"; do
    if [ "$first" = "true" ]; then
      first=false
    else
      WARNINGS_JSON+=","
    fi
    # Escape quotes in message
    escaped_msg=$(echo "$msg" | sed 's/"/\\"/g')
    WARNINGS_JSON+="\"$escaped_msg\""
  done
  WARNINGS_JSON+="]"

  cat > "$SUMMARY_FILE" <<EOF
{
  "status": "completed_with_warnings",
  "applied_at": "$(date -Iseconds)",
  "has_warnings": true,
  "warnings": $WARNINGS_JSON,
  "message": "Local state applied with ${#WARNING_MESSAGES[@]} warning(s) - check git status in container"
}
EOF
else
  cat > "$SUMMARY_FILE" <<EOF
{
  "status": "success",
  "applied_at": "$(date -Iseconds)",
  "has_warnings": false,
  "message": "Local state applied successfully"
}
EOF
fi

exit 0
