#!/bin/bash
set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

echo -e "${GREEN}🔧 Updating IAM Trust Policies for Pipeline Execution Roles${NC}"
echo ""

# Source the .env file for configuration
if [ -f .env ]; then
    export $(cat .env | grep -v '^#' | xargs)
else
    echo -e "${RED}Error: .env file not found${NC}"
    echo "Please run bootstrap-stage1.py first to create .env file"
    exit 1
fi

# Validate required environment variables
if [ -z "$GH_ACCOUNT" ] || [ -z "$GH_REPO" ]; then
    echo -e "${RED}Error: GH_ACCOUNT and GH_REPO must be set in .env${NC}"
    exit 1
fi

# Function to extract role name from ARN
get_role_name() {
    echo "$1" | awk -F'/' '{print $NF}'
}

# Function to get account ID from role ARN
get_account_id() {
    echo "$1" | awk -F':' '{print $5}'
}

# Create trust policy for a role
create_trust_policy() {
    local ACCOUNT_ID=$1
    local REPO_OWNER=$2
    local REPO_NAME=$3

    cat > /tmp/trust-policy-${ACCOUNT_ID}.json <<EOF
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Federated": "arn:aws:iam::${ACCOUNT_ID}:oidc-provider/token.actions.githubusercontent.com"
            },
            "Action": "sts:AssumeRoleWithWebIdentity",
            "Condition": {
                "StringEquals": {
                    "token.actions.githubusercontent.com:aud": "sts.amazonaws.com"
                },
                "StringLike": {
                    "token.actions.githubusercontent.com:sub": [
                        "repo:${REPO_OWNER}/${REPO_NAME}:ref:refs/heads/main",
                        "repo:${REPO_OWNER}/${REPO_NAME}:ref:refs/heads/feat/*",
                        "repo:${REPO_OWNER}/${REPO_NAME}:ref:refs/heads/fix/*",
                        "repo:${REPO_OWNER}/${REPO_NAME}:ref:refs/heads/docs/*",
                        "repo:${REPO_OWNER}/${REPO_NAME}:ref:refs/heads/refactor/*",
                        "repo:${REPO_OWNER}/${REPO_NAME}:ref:refs/heads/test/*",
                        "repo:${REPO_OWNER}/${REPO_NAME}:ref:refs/heads/chore/*",
                        "repo:${REPO_OWNER}/${REPO_NAME}:ref:refs/heads/perf/*",
                        "repo:${REPO_OWNER}/${REPO_NAME}:pull_request"
                    ]
                }
            }
        }
    ]
}
EOF

}

# Function to update trust policy for a role
update_role_trust_policy() {
    local ROLE_ARN=$1
    local ENV_NAME=$2
    local AWS_PROFILE=$3

    if [ -z "$ROLE_ARN" ]; then
        echo -e "${YELLOW}⚠️  No ${ENV_NAME} role ARN found in .env, skipping...${NC}"
        return
    fi

    ROLE_NAME=$(get_role_name "$ROLE_ARN")
    ACCOUNT_ID=$(get_account_id "$ROLE_ARN")

    echo -e "${YELLOW}=== Updating ${ENV_NAME} Pipeline Execution Role ===${NC}"
    echo "Role: $ROLE_NAME"
    echo "Account: $ACCOUNT_ID"
    echo "Profile: $AWS_PROFILE"
    echo ""

    # Create the trust policy
    create_trust_policy "$ACCOUNT_ID" "$GH_ACCOUNT" "$GH_REPO"

    echo "New trust policy:"
    cat /tmp/trust-policy-${ACCOUNT_ID}.json | jq '.'
    echo ""

    # Update the role
    aws iam update-assume-role-policy \
        --role-name "$ROLE_NAME" \
        --policy-document file:///tmp/trust-policy-${ACCOUNT_ID}.json \
        --profile "$AWS_PROFILE" 2>/dev/null

    if [ $? -eq 0 ]; then
        echo -e "${GREEN}✅ Successfully updated ${ENV_NAME} role trust policy${NC}"
    else
        echo -e "${RED}❌ Failed to update ${ENV_NAME} role trust policy${NC}"
        echo "   Make sure the AWS profile '$AWS_PROFILE' is configured and has permissions"
    fi

    # Clean up
    rm -f /tmp/trust-policy-${ACCOUNT_ID}.json

    echo ""
    echo ""
}

# Update roles based on what's available in .env
if [ -n "$TESTING_PIPELINE_EXECUTION_ROLE" ]; then
    update_role_trust_policy "$TESTING_PIPELINE_EXECUTION_ROLE" "Testing" "default"
fi

if [ -n "$STAGING_PIPELINE_EXECUTION_ROLE" ]; then
    update_role_trust_policy "$STAGING_PIPELINE_EXECUTION_ROLE" "Staging" "default"
fi

if [ -n "$PROD_PIPELINE_EXECUTION_ROLE" ]; then
    update_role_trust_policy "$PROD_PIPELINE_EXECUTION_ROLE" "Production" "aillc"
fi

echo -e "${GREEN}🎉 Trust policy updates completed!${NC}"
echo ""
echo "The roles now support the ${GH_REPO} branch strategy:"
echo "  ✅ Push events to: main"
echo "  ✅ Push events to: feat/*, fix/*, docs/*, refactor/*, test/*, chore/*, perf/*"
echo "  ✅ Pull request events (from any branch)"
echo ""
echo -e "${YELLOW}GitHub Actions will now be able to assume roles for:${NC}"
echo "  - Direct pushes to main and feature branches"
echo "  - Pull requests targeting any branch"
