#!/usr/bin/env bash

#########################################################################################
#### lambda-deployer.sh — Package, deploy, and schedule Python Lambda functions      ####
#### Supports dependency bundling, EventBridge scheduling, invocation, and log tail   ####
#### Requires: bash 4+, aws-cli v2, jq, zip, pip3                                   ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   export AWS_PROFILE="production"                                               ####
####   ./lambda-deployer.sh --deploy --function-name my-func --role-arn <ARN>        ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
AWS_REGION="${AWS_REGION:-}"
FUNCTION_NAME="${FUNCTION_NAME:-}"
LAMBDA_RUNTIME="${LAMBDA_RUNTIME:-python3.12}"
LAMBDA_HANDLER="${LAMBDA_HANDLER:-lambda_function.lambda_handler}"
LAMBDA_ROLE_ARN="${LAMBDA_ROLE_ARN:-}"
LAMBDA_TIMEOUT="${LAMBDA_TIMEOUT:-30}"
LAMBDA_MEMORY="${LAMBDA_MEMORY:-128}"
LAMBDA_ENV_VARS="${LAMBDA_ENV_VARS:-}"
LAMBDA_LAYERS="${LAMBDA_LAYERS:-}"
SOURCE_DIR="${SOURCE_DIR:-.}"
SCHEDULE_EXPRESSION="${SCHEDULE_EXPRESSION:-}"
PAYLOAD="${PAYLOAD:-}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
RUN_MODE=""
ZIP_FILE=""
START_TIME=""
TEMP_DIR=""

# ── Colors ────────────────────────────────────────────────────────────
RED="" GREEN="" YELLOW="" BLUE="" BOLD="" RESET=""
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" RESET=""
        return
    fi
    if [[ "$COLOR" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        BOLD='\033[1m'
        RESET='\033[0m'
    else
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" RESET=""
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${BLUE}[DEBUG]${RESET} $*"; fi; }

die() {
    err "$@"
    exit 1
}

# ── Cleanup ───────────────────────────────────────────────────────────
cleanup() {
    if [[ -n "$TEMP_DIR" ]] && [[ -d "$TEMP_DIR" ]]; then
        verbose "Cleaning up temp directory: $TEMP_DIR"
        rm -rf "$TEMP_DIR"
    fi
}

trap cleanup EXIT

# ── AWS CLI wrapper ───────────────────────────────────────────────────
aws_cmd() {
    local args=("$@")
    [[ -n "$AWS_REGION" ]] && args+=(--region "$AWS_REGION")
    verbose "aws ${args[*]}"
    aws "${args[@]}"
}

# ── Resolve region ────────────────────────────────────────────────────
resolve_region() {
    [[ -z "$AWS_REGION" ]] && AWS_REGION="$(aws configure get region 2>/dev/null || true)"
    [[ -z "$AWS_REGION" ]] && die "Cannot determine AWS region. Set AWS_REGION or configure aws-cli."
    verbose "Region: $AWS_REGION"
}

# ── Dependency checks ────────────────────────────────────────────────
check_dependencies() {
    local missing=()
    for cmd in aws jq zip pip3; do
        command -v "$cmd" &>/dev/null || missing+=("$cmd")
    done
    [[ ${#missing[@]} -gt 0 ]] && die "Missing required tools: ${missing[*]}"
    [[ "${BASH_VERSINFO[0]}" -lt 4 ]] && die "Bash 4+ required (found ${BASH_VERSION})"
}

# ── Print header ──────────────────────────────────────────────────────
print_header() {
    echo -e "${BOLD}Lambda Deployer${RESET}"
    echo "Region: $AWS_REGION"
    echo "Mode:   $RUN_MODE"
    echo "Time:   $(date -u +%Y-%m-%dT%H:%M:%SZ)"
    echo ""
}

# ── Elapsed time ──────────────────────────────────────────────────────
elapsed() {
    local end
    end=$(date +%s)
    echo $(( end - START_TIME ))
}

# ── Build environment variables JSON ─────────────────────────────────
build_env_vars_json() {
    local env_str="$1"
    [[ -z "$env_str" ]] && return
    local json="{"
    local first=true
    IFS=',' read -ra pairs <<< "$env_str"
    for pair in "${pairs[@]}"; do
        [[ "$first" == "true" ]] && first=false || json+=","
        json+="\"${pair%%=*}\":\"${pair#*=}\""
    done
    echo "{\"Variables\":${json}}}"
}

# ── Build layers array ───────────────────────────────────────────────
build_layers_args() {
    [[ -z "$1" ]] && return
    echo "${1//,/ }"
}

# ── Package mode ──────────────────────────────────────────────────────
do_package() {
    [[ -z "$FUNCTION_NAME" ]] && die "--function-name is required for package mode"

    local src_dir
    src_dir="$(realpath "$SOURCE_DIR")"
    [[ -d "$src_dir" ]] || die "Source directory not found: $src_dir"

    log "Packaging function ${BOLD}$FUNCTION_NAME${RESET}..."

    TEMP_DIR="$(mktemp -d)"
    local pkg_dir="$TEMP_DIR/package"
    mkdir -p "$pkg_dir"

    # Install dependencies
    local req_file="$src_dir/requirements.txt"
    if [[ -f "$req_file" ]]; then
        log "Installing dependencies from requirements.txt..."
        pip3 install -r "$req_file" -t "$pkg_dir" --quiet --disable-pip-version-check 2>/dev/null \
            || die "pip3 install failed"

        # Clean up pip metadata to reduce zip size
        find "$pkg_dir" -type d \( -name "__pycache__" -o -name "*.dist-info" -o -name "*.egg-info" \) \
            -exec rm -rf {} + 2>/dev/null || true

        local pkg_count
        pkg_count=$(grep -cE '^[^#[:space:]]' "$req_file" || echo "0")
        log "Collected $pkg_count package(s)"
    else
        warn "No requirements.txt found — packaging handler only"
    fi

    # Copy handler code
    local py_count=0
    while IFS= read -r -d '' f; do
        cp "$f" "$pkg_dir/"
        ((py_count++)) || true
    done < <(find "$src_dir" -maxdepth 1 -name "*.py" -print0)

    if [[ $py_count -eq 0 ]]; then
        die "No .py files found in $src_dir"
    fi
    verbose "Copied $py_count Python file(s)"

    # Create zip
    ZIP_FILE="/tmp/lambda-${FUNCTION_NAME}.zip"
    (cd "$pkg_dir" && zip -r -q "$ZIP_FILE" .) \
        || die "Failed to create zip"

    local size
    size=$(du -h "$ZIP_FILE" | cut -f1)
    log "Created deployment package: ${BOLD}$ZIP_FILE${RESET} ($size)"
}

# ── Check if function exists ─────────────────────────────────────────
function_exists() {
    local name="$1"
    aws_cmd lambda get-function --function-name "$name" &>/dev/null
}

# ── Deploy mode ───────────────────────────────────────────────────────
do_deploy() {
    [[ -z "$FUNCTION_NAME" ]] && die "--function-name is required for deploy mode"

    # Package first if zip doesn't exist
    ZIP_FILE="/tmp/lambda-${FUNCTION_NAME}.zip"
    if [[ ! -f "$ZIP_FILE" ]]; then
        do_package
    fi

    if function_exists "$FUNCTION_NAME"; then
        log "Function ${BOLD}$FUNCTION_NAME${RESET} exists — updating..."

        aws_cmd lambda update-function-code \
            --function-name "$FUNCTION_NAME" \
            --zip-file "fileb://$ZIP_FILE" \
            --output text --query 'FunctionArn' >/dev/null \
            || die "Failed to update function code"

        aws_cmd lambda wait function-updated \
            --function-name "$FUNCTION_NAME" 2>/dev/null || true

        local config_args=(
            lambda update-function-configuration
            --function-name "$FUNCTION_NAME"
            --runtime "$LAMBDA_RUNTIME"
            --handler "$LAMBDA_HANDLER"
            --timeout "$LAMBDA_TIMEOUT"
            --memory-size "$LAMBDA_MEMORY"
        )

        if [[ -n "$LAMBDA_ENV_VARS" ]]; then
            local env_json
            env_json=$(build_env_vars_json "$LAMBDA_ENV_VARS")
            config_args+=(--environment "$env_json")
        fi

        if [[ -n "$LAMBDA_LAYERS" ]]; then
            local layers
            layers=$(build_layers_args "$LAMBDA_LAYERS")
            # shellcheck disable=SC2086,SC2206
            config_args+=(--layers $layers)
        fi

        local fn_arn
        fn_arn=$(aws_cmd "${config_args[@]}" --output text --query 'FunctionArn') \
            || die "Failed to update function configuration"

        echo -e "  ${GREEN}✓${RESET} Function updated: $fn_arn"
    else
        [[ -z "$LAMBDA_ROLE_ARN" ]] && die "--role-arn is required to create a new function"

        log "Function ${BOLD}$FUNCTION_NAME${RESET} does not exist — creating..."
        log "Creating function $FUNCTION_NAME ($LAMBDA_RUNTIME, $LAMBDA_MEMORY MB, ${LAMBDA_TIMEOUT}s timeout)"

        local create_args=(
            lambda create-function
            --function-name "$FUNCTION_NAME"
            --runtime "$LAMBDA_RUNTIME"
            --handler "$LAMBDA_HANDLER"
            --role "$LAMBDA_ROLE_ARN"
            --timeout "$LAMBDA_TIMEOUT"
            --memory-size "$LAMBDA_MEMORY"
            --zip-file "fileb://$ZIP_FILE"
        )

        if [[ -n "$LAMBDA_ENV_VARS" ]]; then
            local env_json
            env_json=$(build_env_vars_json "$LAMBDA_ENV_VARS")
            create_args+=(--environment "$env_json")
        fi

        if [[ -n "$LAMBDA_LAYERS" ]]; then
            local layers
            layers=$(build_layers_args "$LAMBDA_LAYERS")
            # shellcheck disable=SC2086,SC2206
            create_args+=(--layers $layers)
        fi

        local fn_arn
        fn_arn=$(aws_cmd "${create_args[@]}" --output text --query 'FunctionArn') \
            || die "Failed to create function"

        # Wait for function to become active
        verbose "Waiting for function to become active..."
        aws_cmd lambda wait function-active-v2 \
            --function-name "$FUNCTION_NAME" 2>/dev/null || true

        echo -e "  ${GREEN}✓${RESET} Function created: $fn_arn"
    fi
}

# ── Schedule mode ─────────────────────────────────────────────────────
do_schedule() {
    [[ -z "$FUNCTION_NAME" ]] && die "--function-name is required for schedule mode"
    [[ -z "$SCHEDULE_EXPRESSION" ]] && die "--schedule-expression is required for schedule mode"

    log "Configuring EventBridge schedule for ${BOLD}$FUNCTION_NAME${RESET}..."

    local fn_arn
    fn_arn=$(aws_cmd lambda get-function \
        --function-name "$FUNCTION_NAME" \
        --output text --query 'Configuration.FunctionArn' 2>/dev/null) \
        || die "Function $FUNCTION_NAME not found — deploy it first"

    local rule_name="lambda-deployer-${FUNCTION_NAME}"
    local rule_arn
    rule_arn=$(aws_cmd events put-rule \
        --name "$rule_name" \
        --schedule-expression "$SCHEDULE_EXPRESSION" \
        --state ENABLED \
        --description "Scheduled trigger for $FUNCTION_NAME (managed by lambda-deployer)" \
        --output text --query 'RuleArn') \
        || die "Failed to create EventBridge rule"

    echo -e "  ${GREEN}✓${RESET} Rule: $rule_name ($SCHEDULE_EXPRESSION)"

    aws_cmd events put-targets \
        --rule "$rule_name" \
        --targets "[{\"Id\":\"${FUNCTION_NAME}-target\",\"Arn\":\"${fn_arn}\"}]" \
        --output text >/dev/null \
        || die "Failed to add Lambda target to rule"

    local stmt_id="lambda-deployer-${FUNCTION_NAME}-invoke"
    aws_cmd lambda remove-permission \
        --function-name "$FUNCTION_NAME" \
        --statement-id "$stmt_id" 2>/dev/null || true

    aws_cmd lambda add-permission \
        --function-name "$FUNCTION_NAME" \
        --statement-id "$stmt_id" \
        --action "lambda:InvokeFunction" \
        --principal "events.amazonaws.com" \
        --source-arn "$rule_arn" \
        --output text >/dev/null \
        || die "Failed to add invoke permission"

    echo -e "  ${GREEN}✓${RESET} Permission granted for EventBridge to invoke $FUNCTION_NAME"
    log "Schedule configured: $SCHEDULE_EXPRESSION"
}

# ── Invoke mode ───────────────────────────────────────────────────────
do_invoke() {
    [[ -z "$FUNCTION_NAME" ]] && die "--function-name is required for invoke mode"

    log "Invoking ${BOLD}$FUNCTION_NAME${RESET}..."

    local invoke_args=(
        lambda invoke
        --function-name "$FUNCTION_NAME"
        --log-type Tail
    )

    if [[ -n "$PAYLOAD" ]]; then
        invoke_args+=(--payload "$PAYLOAD")
        verbose "Payload: $PAYLOAD"
    fi

    local output_file
    output_file="$(mktemp)"

    local response
    response=$(aws_cmd "${invoke_args[@]}" "$output_file" 2>&1) \
        || die "Invoke failed: $response"

    local func_error
    func_error=$(echo "$response" | jq -r '.FunctionError // empty' 2>/dev/null || echo "")

    if [[ -n "$func_error" ]]; then
        echo -e "  ${RED}✗${RESET} Function error: $func_error"
    else
        echo -e "  ${GREEN}✓${RESET} Status: $(echo "$response" | jq -r '.StatusCode // 200' 2>/dev/null)"
    fi

    echo ""
    echo -e "${BOLD}Response:${RESET}"
    jq '.' "$output_file" 2>/dev/null || cat "$output_file"
    echo ""

    local log_result
    log_result=$(echo "$response" | jq -r '.LogResult // empty' 2>/dev/null || echo "")
    if [[ -n "$log_result" ]]; then
        echo -e "${BOLD}Execution Log:${RESET}"
        echo "$log_result" | base64 --decode 2>/dev/null || true
    fi

    rm -f "$output_file"
}

# ── Logs mode ─────────────────────────────────────────────────────────
do_logs() {
    [[ -z "$FUNCTION_NAME" ]] && die "--function-name is required for logs mode"

    local log_group="/aws/lambda/$FUNCTION_NAME"
    local exists
    exists=$(aws_cmd logs describe-log-groups \
        --log-group-name-prefix "$log_group" \
        --query "logGroups[?logGroupName=='$log_group'].logGroupName" \
        --output text 2>/dev/null || echo "")
    [[ -z "$exists" ]] && die "Log group $log_group not found — has $FUNCTION_NAME been invoked?"

    log "Tailing logs for ${BOLD}$FUNCTION_NAME${RESET} (Ctrl+C to stop)..."
    echo ""

    local start_time next_token=""
    start_time=$(( $(date +%s) * 1000 - 300000 ))

    while true; do
        local filter_args=(logs filter-log-events --log-group-name "$log_group"
            --start-time "$start_time" --interleaved)
        [[ -n "$next_token" ]] && filter_args+=(--next-token "$next_token")

        local result
        result=$(aws_cmd "${filter_args[@]}" --output json 2>/dev/null || echo "{}")

        local events
        events=$(echo "$result" | jq -r \
            '.events[]? | "\(.timestamp | . / 1000 | strftime("%Y-%m-%dT%H:%M:%SZ")) \(.message)"' \
            2>/dev/null || echo "")
        [[ -n "$events" ]] && echo "$events"

        local new_token
        new_token=$(echo "$result" | jq -r '.nextToken // empty' 2>/dev/null || echo "")
        if [[ -n "$new_token" ]]; then
            next_token="$new_token"
        else
            next_token=""
            local last_ts
            last_ts=$(echo "$result" | jq -r '.events[-1]?.timestamp // empty' 2>/dev/null || echo "")
            [[ -n "$last_ts" ]] && start_time=$(( last_ts + 1 ))
        fi
        sleep 2
    done
}

# ── List mode ─────────────────────────────────────────────────────────
do_list() {
    log "Listing Lambda functions in ${BOLD}$AWS_REGION${RESET}..."
    echo ""

    local functions
    functions=$(aws_cmd lambda list-functions \
        --query 'Functions[*].[FunctionName,Runtime,CodeSize,LastModified]' \
        --output json 2>/dev/null) \
        || die "Failed to list functions"

    local count
    count=$(echo "$functions" | jq 'length' 2>/dev/null || echo "0")

    if [[ "$count" -eq 0 ]]; then
        log "No Lambda functions found in $AWS_REGION"
        return
    fi

    printf "  ${BOLD}%-25s %-14s %-10s %s${RESET}\n" "FUNCTION" "RUNTIME" "SIZE" "LAST MODIFIED"
    echo "  ─────────────────────────────────────────────────────────────────"

    echo "$functions" | jq -r '.[] | @tsv' 2>/dev/null | while IFS=$'\t' read -r name runtime size modified; do
        # Convert bytes to human-readable
        local human_size
        if [[ "$size" -ge 1048576 ]]; then
            human_size="$(awk "BEGIN{printf \"%.1f MB\", $size/1048576}")"
        elif [[ "$size" -ge 1024 ]]; then
            human_size="$(awk "BEGIN{printf \"%.1f KB\", $size/1024}")"
        else
            human_size="${size} B"
        fi

        # Trim the modified timestamp
        local short_modified
        short_modified="${modified%%+*}"
        short_modified="${short_modified%.*}"

        printf "  %-25s %-14s %-10s %s\n" "$name" "$runtime" "$human_size" "$short_modified"
    done

    echo ""
    log "Total: $count function(s)"
}

# ── Usage / help ──────────────────────────────────────────────────────
usage() {
    cat <<EOF
Usage: $SCRIPT_NAME <mode> [options]

Modes:
  --package    Package handler + deps into zip   --deploy   Create/update function
  --schedule   EventBridge scheduled rule        --invoke   Invoke function
  --logs       Tail CloudWatch Logs              --list     List functions

Options:
  --function-name NAME       Function name       --runtime RUNTIME   (python3.12)
  --handler HANDLER          Entry point          --role-arn ARN      Exec role
  --timeout N                Seconds (30)         --memory N          MB (128)
  --env-vars K=V,...         Env variables        --layers ARN,...    Layer ARNs
  --schedule-expression E    Cron/rate expr       --payload JSON      Invoke payload
  --source-dir DIR           Source dir (.)       --verbose           Debug output
  --no-color                 No ANSI colors       -h, --help          This help
EOF
}

# ── Parse arguments ───────────────────────────────────────────────────
parse_args() {
    if [[ $# -eq 0 ]]; then
        usage
        exit 0
    fi

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --package|--deploy|--schedule|--invoke|--logs|--list)
                RUN_MODE="${1#--}"; shift ;;
            --function-name)        [[ $# -lt 2 ]] && die "$1 requires a value"; FUNCTION_NAME="$2"; shift 2 ;;
            --runtime)              [[ $# -lt 2 ]] && die "$1 requires a value"; LAMBDA_RUNTIME="$2"; shift 2 ;;
            --handler)              [[ $# -lt 2 ]] && die "$1 requires a value"; LAMBDA_HANDLER="$2"; shift 2 ;;
            --role-arn)             [[ $# -lt 2 ]] && die "$1 requires a value"; LAMBDA_ROLE_ARN="$2"; shift 2 ;;
            --timeout)              [[ $# -lt 2 ]] && die "$1 requires a value"; LAMBDA_TIMEOUT="$2"; shift 2 ;;
            --memory)               [[ $# -lt 2 ]] && die "$1 requires a value"; LAMBDA_MEMORY="$2"; shift 2 ;;
            --env-vars)             [[ $# -lt 2 ]] && die "$1 requires a value"; LAMBDA_ENV_VARS="$2"; shift 2 ;;
            --layers)               [[ $# -lt 2 ]] && die "$1 requires a value"; LAMBDA_LAYERS="$2"; shift 2 ;;
            --schedule-expression)  [[ $# -lt 2 ]] && die "$1 requires a value"; SCHEDULE_EXPRESSION="$2"; shift 2 ;;
            --payload)              [[ $# -lt 2 ]] && die "$1 requires a value"; PAYLOAD="$2"; shift 2 ;;
            --source-dir)           [[ $# -lt 2 ]] && die "$1 requires a value"; SOURCE_DIR="$2"; shift 2 ;;
            --verbose)              VERBOSE="true"; shift ;;
            --no-color)             COLOR="never"; shift ;;
            -h|--help)              usage; exit 0 ;;
            *)                      die "Unknown option: $1 (see --help)" ;;
        esac
    done

    if [[ -z "$RUN_MODE" ]]; then err "No mode specified"; echo ""; usage; exit 1; fi
}

# ── Main ──────────────────────────────────────────────────────────────
main() {
    parse_args "$@"
    setup_colors
    check_dependencies
    resolve_region

    START_TIME=$(date +%s)
    print_header

    case "$RUN_MODE" in
        package)
            do_package
            ;;
        deploy)
            do_deploy
            ;;
        schedule)
            do_schedule
            ;;
        invoke)
            do_invoke
            ;;
        logs)
            do_logs
            ;;
        list)
            do_list
            ;;
        *)
            die "Unknown mode: $RUN_MODE"
            ;;
    esac

    if [[ "$RUN_MODE" != "logs" ]]; then
        log "Completed in $(elapsed)s"
    fi
}

main "$@"