#!/usr/bin/env python3
"""Patches all deploy-services.sh files to add retry logic for helm s3 push.
Fixes PreconditionFailed 412 race condition when concurrent pipelines update S3 index.yaml."""

import os
import re
import sys

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

RETRY_FUNC = '''
# Retry helm s3 push with backoff to handle S3 PreconditionFailed 412 race conditions
helm_s3_push_with_retry() {
    local chart_file="$1"
    local repo_name="$2"
    local max_retries=5
    local attempt=1
    local wait_time=5
    while [ $attempt -le $max_retries ]; do
        if helm s3 push "${chart_file}" "${repo_name}" --force 2>&1; then
            return 0
        fi
        echo "helm s3 push failed (attempt ${attempt}/${max_retries}), retrying in ${wait_time}s..."
        sleep $wait_time
        wait_time=$((wait_time * 2))
        attempt=$((attempt + 1))
    done
    echo "ERROR: helm s3 push failed after ${max_retries} attempts"
    return 1
}
'''

patched = 0
skipped = 0

for root, dirs, files in os.walk(REPO_ROOT):
    for f in files:
        if f != "deploy-services.sh":
            continue
        filepath = os.path.join(root, f)
        # Skip this patch script itself
        if "scripts/patch-" in filepath:
            continue

        with open(filepath, "r") as fh:
            content = fh.read()

        if "helm_s3_push_with_retry" in content:
            print(f"SKIP: {os.path.relpath(filepath, REPO_ROOT)}")
            skipped += 1
            continue

        if "helm s3 push" not in content:
            print(f"SKIP (no push): {os.path.relpath(filepath, REPO_ROOT)}")
            skipped += 1
            continue

        # Replace bare helm s3 push with retry call, capturing the repo name
        content = re.sub(
            r'helm s3 push \$\{chart_file\} (tbi-charts\S*) --force',
            r'helm_s3_push_with_retry ${chart_file} \1',
            content,
        )

        # Inject retry function after 'set -e'
        content = content.replace("set -e\n", "set -e\n" + RETRY_FUNC, 1)

        with open(filepath, "w") as fh:
            fh.write(content)

        patched += 1
        print(f"PATCHED: {os.path.relpath(filepath, REPO_ROOT)}")

print(f"\n=== Summary ===")
print(f"Patched: {patched}")
print(f"Skipped: {skipped}")