#!/usr/bin/env bash

#########################################################################################
#### ebs-snapshot-manager.sh — Create, manage, audit, and prune AWS EBS snapshots    ####
#### Supports automated creation, cross-region copy, retention, and orphan detection ####
#### Requires: bash 4+, aws-cli v2, jq                                              ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   export AWS_PROFILE="production"                                               ####
####   ./ebs-snapshot-manager.sh --snapshot                                           ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
AWS_REGION="${AWS_REGION:-}"
VOLUME_IDS="${VOLUME_IDS:-}"
VOLUME_TAG_KEY="${VOLUME_TAG_KEY:-}"
VOLUME_TAG_VALUE="${VOLUME_TAG_VALUE:-}"
RETENTION_DAYS="${RETENTION_DAYS:-30}"
COPY_TO_REGION="${COPY_TO_REGION:-}"
SNAPSHOT_DESCRIPTION="${SNAPSHOT_DESCRIPTION:-Automated snapshot by ebs-snapshot-manager}"
NO_WAIT="${NO_WAIT:-false}"
DRY_RUN="${DRY_RUN:-true}"
RESTORE_AZ="${RESTORE_AZ:-}"
RESTORE_VOLUME_TYPE="${RESTORE_VOLUME_TYPE:-gp3}"
RESTORE_IOPS="${RESTORE_IOPS:-}"
RESTORE_THROUGHPUT="${RESTORE_THROUGHPUT:-}"
OUTPUT_FORMAT="${OUTPUT_FORMAT:-text}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
RUN_MODE=""
TARGET_VOLUME=""
TARGET_SNAPSHOT=""
START_TIME=""
WARNINGS=0

# ── Colors ────────────────────────────────────────────────────────────
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; ((WARNINGS++)) || true; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${BLUE}[DEBUG]${RESET} $*"; fi; }

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

# ── Dependency check ──────────────────────────────────────────────────
check_deps() {
    for cmd in aws jq; do
        if ! command -v "$cmd" &>/dev/null; then
            err "${cmd} is required but not installed"
            exit 1
        fi
    done

    # Verify AWS credentials
    if ! aws sts get-caller-identity &>/dev/null; then
        err "AWS credentials not configured or expired"
        exit 1
    fi

    # Determine region
    if [[ -z "$AWS_REGION" ]]; then
        AWS_REGION=$(aws configure get region 2>/dev/null || echo "")
        if [[ -z "$AWS_REGION" ]]; then
            err "AWS_REGION is required (set via env var or aws configure)"
            exit 1
        fi
    fi

    verbose "Using region: ${AWS_REGION}"
    verbose "Account: $(aws sts get-caller-identity --query 'Account' --output text 2>/dev/null)"
}

# ── Get volume list ───────────────────────────────────────────────────
get_volumes() {
    local filters=()

    if [[ -n "$VOLUME_IDS" ]]; then
        # Specific volumes requested
        local vol_array
        IFS=',' read -ra vol_array <<< "$VOLUME_IDS"
        aws_cmd ec2 describe-volumes \
            --volume-ids "${vol_array[@]}" \
            --query 'Volumes[*].VolumeId' \
            --output text | tr '\t' '\n'
        return
    fi

    if [[ -n "$VOLUME_TAG_KEY" ]]; then
        filters+=(--filters "Name=tag:${VOLUME_TAG_KEY},Values=${VOLUME_TAG_VALUE:-*}")
    fi

    aws_cmd ec2 describe-volumes \
        "${filters[@]}" \
        --query 'Volumes[*].VolumeId' \
        --output text | tr '\t' '\n'
}

# ── Get account ID ───────────────────────────────────────────────────
get_account_id() {
    aws sts get-caller-identity --query 'Account' --output text
}

# ══════════════════════════════════════════════════════════════════════
#                          SNAPSHOT MODE
# ══════════════════════════════════════════════════════════════════════

do_snapshot() {
    log "Creating EBS snapshots..."
    local volumes
    volumes=$(get_volumes)

    if [[ -z "$volumes" ]]; then
        warn "No volumes found matching criteria"
        return
    fi

    local vol_count
    vol_count=$(echo "$volumes" | wc -l)
    log "Found ${vol_count} volume(s) to snapshot"

    local created=0
    local failed=0
    local snapshot_ids=()
    local now
    now=$(date -u +%Y-%m-%dT%H:%M:%SZ)

    while IFS= read -r vol_id; do
        [[ -z "$vol_id" ]] && continue
        verbose "Snapshotting ${vol_id}..."

        local vol_name
        # shellcheck disable=SC2016
        vol_name=$(aws_cmd ec2 describe-volumes \
            --volume-ids "$vol_id" \
            --query 'Volumes[0].Tags[?Key==`Name`].Value | [0]' \
            --output text 2>/dev/null) || vol_name="N/A"
        [[ "$vol_name" == "None" ]] && vol_name="N/A"

        local snap_id
        snap_id=$(aws_cmd ec2 create-snapshot \
            --volume-id "$vol_id" \
            --description "$SNAPSHOT_DESCRIPTION" \
            --tag-specifications "ResourceType=snapshot,Tags=[
                {Key=Name,Value=snap-${vol_id}-$(date +%Y%m%d)},
                {Key=CreatedBy,Value=ebs-snapshot-manager},
                {Key=CreatedAt,Value=${now}},
                {Key=VolumeId,Value=${vol_id}},
                {Key=VolumeName,Value=${vol_name}}
            ]" \
            --query 'SnapshotId' \
            --output text 2>/dev/null) || snap_id=""

        if [[ -n "$snap_id" ]]; then
            echo -e "  ${GREEN}✓${RESET} ${vol_id} → ${snap_id} (${vol_name})"
            snapshot_ids+=("$snap_id")
            ((created++)) || true
        else
            echo -e "  ${RED}✗${RESET} ${vol_id} — snapshot creation failed"
            ((failed++)) || true
        fi
    done <<< "$volumes"

    # Wait for completion
    if [[ "$NO_WAIT" != "true" && ${#snapshot_ids[@]} -gt 0 ]]; then
        log "Waiting for ${#snapshot_ids[@]} snapshot(s) to complete..."
        for snap_id in "${snapshot_ids[@]}"; do
            if aws_cmd ec2 wait snapshot-completed --snapshot-ids "$snap_id" 2>/dev/null; then
                local size
                size=$(aws_cmd ec2 describe-snapshots \
                    --snapshot-ids "$snap_id" \
                    --query 'Snapshots[0].VolumeSize' \
                    --output text 2>/dev/null) || size="?"
                verbose "${snap_id} completed (${size} GiB)"
            else
                warn "${snap_id} did not complete within timeout"
            fi
        done
    fi

    echo ""
    log "Snapshots created: ${created}, failed: ${failed}"
}

# ══════════════════════════════════════════════════════════════════════
#                          PRUNE MODE
# ══════════════════════════════════════════════════════════════════════

do_prune() {
    local cutoff_epoch
    cutoff_epoch=$(date -d "-${RETENTION_DAYS} days" +%s 2>/dev/null) || \
    cutoff_epoch=$(date -v-"${RETENTION_DAYS}"d +%s 2>/dev/null) || {
        err "Could not calculate retention cutoff date"
        exit 1
    }

    local cutoff_date
    cutoff_date=$(date -d "@${cutoff_epoch}" +%Y-%m-%dT%H:%M:%S 2>/dev/null) || \
    cutoff_date=$(date -r "${cutoff_epoch}" +%Y-%m-%dT%H:%M:%S 2>/dev/null)

    log "Pruning snapshots older than ${RETENTION_DAYS} days (before ${cutoff_date})"
    if [[ "$DRY_RUN" == "true" ]]; then
        log "${YELLOW}DRY RUN${RESET} — no snapshots will be deleted. Use --force to delete."
    fi

    local owner_id
    owner_id=$(get_account_id)

    local snapshots_json
    snapshots_json=$(aws_cmd ec2 describe-snapshots \
        --owner-ids "$owner_id" \
        --filters "Name=tag:CreatedBy,Values=ebs-snapshot-manager" \
        --query 'Snapshots[*].{Id:SnapshotId,Start:StartTime,Size:VolumeSize,Vol:VolumeId}' \
        --output json)

    local total
    total=$(echo "$snapshots_json" | jq 'length')
    echo "$snapshots_json" | jq -c '.[]' | while IFS= read -r snap; do
        local snap_id start_time size vol_id
        snap_id=$(echo "$snap" | jq -r '.Id')
        start_time=$(echo "$snap" | jq -r '.Start')
        size=$(echo "$snap" | jq -r '.Size')
        vol_id=$(echo "$snap" | jq -r '.Vol')

        local snap_epoch
        snap_epoch=$(date -d "$start_time" +%s 2>/dev/null) || \
        snap_epoch=$(date -jf "%Y-%m-%dT%H:%M:%S" "${start_time%%.*}" +%s 2>/dev/null) || snap_epoch=0

        if [[ $snap_epoch -lt $cutoff_epoch ]]; then
            local age_days=$(( ($(date +%s) - snap_epoch) / 86400 ))

            if [[ "$DRY_RUN" == "true" ]]; then
                echo -e "  ${YELLOW}⊘${RESET} ${snap_id} — ${age_days}d old, ${size} GiB, vol: ${vol_id} (would delete)"
            else
                if aws_cmd ec2 delete-snapshot --snapshot-id "$snap_id" 2>/dev/null; then
                    echo -e "  ${GREEN}✓${RESET} ${snap_id} — deleted (${age_days}d old, ${size} GiB)"
                else
                    echo -e "  ${RED}✗${RESET} ${snap_id} — delete failed"
                fi
            fi
        fi
    done

    log "Total managed snapshots: ${total}"
}

# ══════════════════════════════════════════════════════════════════════
#                       COPY-REGION MODE
# ══════════════════════════════════════════════════════════════════════

do_copy_region() {
    if [[ -z "$COPY_TO_REGION" ]]; then
        err "Target region required. Use --copy-region REGION or set COPY_TO_REGION"
        exit 1
    fi

    log "Copying latest snapshots to ${COPY_TO_REGION}..."

    local owner_id
    owner_id=$(get_account_id)

    # Get volumes to copy snapshots for
    local volumes
    if [[ -n "$TARGET_VOLUME" ]]; then
        volumes="$TARGET_VOLUME"
    else
        volumes=$(get_volumes)
    fi

    if [[ -z "$volumes" ]]; then
        warn "No volumes found"
        return
    fi

    local copied=0
    local failed=0

    while IFS= read -r vol_id; do
        [[ -z "$vol_id" ]] && continue

        # Find latest snapshot for this volume
        local latest_snap
        latest_snap=$(aws_cmd ec2 describe-snapshots \
            --owner-ids "$owner_id" \
            --filters "Name=volume-id,Values=${vol_id}" "Name=status,Values=completed" \
            --query 'sort_by(Snapshots, &StartTime)[-1].SnapshotId' \
            --output text 2>/dev/null) || latest_snap=""

        if [[ -z "$latest_snap" || "$latest_snap" == "None" ]]; then
            echo -e "  ${YELLOW}⊘${RESET} ${vol_id} — no completed snapshots found"
            continue
        fi

        # Copy to target region
        local copy_id
        copy_id=$(aws ec2 copy-snapshot \
            --region "$COPY_TO_REGION" \
            --source-region "$AWS_REGION" \
            --source-snapshot-id "$latest_snap" \
            --description "DR copy of ${latest_snap} from ${AWS_REGION}" \
            --tag-specifications "ResourceType=snapshot,Tags=[
                {Key=Name,Value=dr-copy-${latest_snap}},
                {Key=CreatedBy,Value=ebs-snapshot-manager},
                {Key=SourceRegion,Value=${AWS_REGION}},
                {Key=SourceSnapshotId,Value=${latest_snap}},
                {Key=VolumeId,Value=${vol_id}}
            ]" \
            --query 'SnapshotId' \
            --output text 2>/dev/null) || copy_id=""

        if [[ -n "$copy_id" ]]; then
            echo -e "  ${GREEN}✓${RESET} ${latest_snap} → ${copy_id} (${AWS_REGION} → ${COPY_TO_REGION})"
            ((copied++)) || true
        else
            echo -e "  ${RED}✗${RESET} ${latest_snap} — copy failed"
            ((failed++)) || true
        fi
    done <<< "$volumes"

    echo ""
    log "Copied: ${copied}, failed: ${failed}"
}

# ══════════════════════════════════════════════════════════════════════
#                          AUDIT MODE
# ══════════════════════════════════════════════════════════════════════

do_audit() {
    log "Auditing EBS snapshots in ${AWS_REGION}..."

    local owner_id
    owner_id=$(get_account_id)

    local snapshots_json
    snapshots_json=$(aws_cmd ec2 describe-snapshots \
        --owner-ids "$owner_id" \
        --query 'Snapshots[*].{Id:SnapshotId,Vol:VolumeId,Size:VolumeSize,Status:State,Start:StartTime,Desc:Description,Tags:Tags}' \
        --output json)

    local total
    total=$(echo "$snapshots_json" | jq 'length')

    if [[ "$total" -eq 0 ]]; then
        log "No snapshots found"
        return
    fi

    # Get existing volumes for orphan detection
    local existing_volumes
    existing_volumes=$(aws_cmd ec2 describe-volumes \
        --query 'Volumes[*].VolumeId' \
        --output text | tr '\t' '\n' | sort)

    local orphan_count=0
    local untagged_count=0
    local managed_count=0

    echo ""
    echo -e "${BOLD}Snapshot Inventory${RESET}"
    printf "  %-24s %-14s %8s %6s %s\n" "SNAPSHOT" "VOLUME" "SIZE" "AGE" "STATUS"
    echo "  $(printf '%.0s─' {1..70})"

    echo "$snapshots_json" | jq -c '.[]' | while IFS= read -r snap; do
        local snap_id vol_id size status start_time
        snap_id=$(echo "$snap" | jq -r '.Id')
        vol_id=$(echo "$snap" | jq -r '.Vol')
        size=$(echo "$snap" | jq -r '.Size')
        status=$(echo "$snap" | jq -r '.Status')
        start_time=$(echo "$snap" | jq -r '.Start')

        local snap_epoch
        snap_epoch=$(date -d "$start_time" +%s 2>/dev/null) || snap_epoch=0
        local age_days=$(( ($(date +%s) - snap_epoch) / 86400 ))

        # Check if managed
        local is_managed
        is_managed=$(echo "$snap" | jq -r '.Tags // [] | map(select(.Key == "CreatedBy" and .Value == "ebs-snapshot-manager")) | length')

        # Check if orphaned
        local is_orphan="no"
        if ! echo "$existing_volumes" | grep -q "^${vol_id}$" 2>/dev/null; then
            is_orphan="yes"
        fi

        # Check if tagged
        local tag_count
        tag_count=$(echo "$snap" | jq '.Tags // [] | length')

        local status_marker=""
        if [[ "$is_orphan" == "yes" ]]; then
            status_marker="${RED}orphan${RESET}"
        elif [[ "$tag_count" -eq 0 ]]; then
            status_marker="${YELLOW}untagged${RESET}"
        elif [[ "$is_managed" -gt 0 ]]; then
            status_marker="${GREEN}managed${RESET}"
        else
            status_marker="${status}"
        fi

        printf "  %-24s %-14s %6s G %4sd %b\n" \
            "$snap_id" "$vol_id" "$size" "$age_days" "$status_marker"
    done

    # Summary stats
    local total_size
    total_size=$(echo "$snapshots_json" | jq '[.[].Size] | add // 0')
    orphan_count=$(echo "$snapshots_json" | jq --arg vols "$existing_volumes" '
        [.[] | select(.Vol as $v | ($vols | split("\n") | map(select(. != "")) | index($v) == null))] | length
    ')
    untagged_count=$(echo "$snapshots_json" | jq '[.[] | select((.Tags // []) | length == 0)] | length')
    managed_count=$(echo "$snapshots_json" | jq '[.[] | select((.Tags // []) | map(select(.Key == "CreatedBy" and .Value == "ebs-snapshot-manager")) | length > 0)] | length')

    local monthly_cost
    monthly_cost=$(echo "$total_size * 0.05" | bc 2>/dev/null || echo "?")

    echo ""
    echo -e "${BOLD}Summary${RESET}"
    echo -e "  Total snapshots:    ${total}"
    echo -e "  Managed snapshots:  ${managed_count}"
    echo -e "  Total storage:      ${total_size} GiB"
    echo -e "  Est. monthly cost:  \$${monthly_cost}"
    echo -e "  Orphaned:           ${orphan_count}"
    echo -e "  Untagged:           ${untagged_count}"

    if [[ "$orphan_count" -gt 0 ]]; then
        echo ""
        warn "${orphan_count} orphaned snapshot(s) found — source volume no longer exists"
    fi

    # Check volumes without recent snapshots
    echo ""
    echo -e "${BOLD}Volumes Without Recent Snapshots (>${RETENTION_DAYS}d)${RESET}"

    local volumes_json
    volumes_json=$(aws_cmd ec2 describe-volumes \
        --query 'Volumes[*].{Id:VolumeId,Size:Size,State:State}' \
        --output json)

    echo "$volumes_json" | jq -c '.[]' | while IFS= read -r vol; do
        local v_id v_size
        v_id=$(echo "$vol" | jq -r '.Id')
        v_size=$(echo "$vol" | jq -r '.Size')

        local latest_snap_time
        latest_snap_time=$(aws_cmd ec2 describe-snapshots \
            --owner-ids "$owner_id" \
            --filters "Name=volume-id,Values=${v_id}" "Name=status,Values=completed" \
            --query 'sort_by(Snapshots, &StartTime)[-1].StartTime' \
            --output text 2>/dev/null) || latest_snap_time="None"

        if [[ "$latest_snap_time" == "None" || -z "$latest_snap_time" ]]; then
            echo -e "  ${RED}✗${RESET} ${v_id} (${v_size} GiB) — ${RED}no snapshots${RESET}"
        else
            local snap_epoch
            snap_epoch=$(date -d "$latest_snap_time" +%s 2>/dev/null) || snap_epoch=0
            local cutoff_epoch
            cutoff_epoch=$(date -d "-${RETENTION_DAYS} days" +%s 2>/dev/null) || cutoff_epoch=0

            if [[ $snap_epoch -lt $cutoff_epoch ]]; then
                local age=$(( ($(date +%s) - snap_epoch) / 86400 ))
                echo -e "  ${YELLOW}!${RESET} ${v_id} (${v_size} GiB) — last snapshot ${age}d ago"
            fi
        fi
    done

    # Prometheus output
    if [[ "$OUTPUT_FORMAT" == "prometheus" ]]; then
        echo ""
        echo "# HELP ebs_snapshots_total Total EBS snapshots"
        echo "# TYPE ebs_snapshots_total gauge"
        echo "ebs_snapshots_total{region=\"${AWS_REGION}\"} ${total}"
        echo "# HELP ebs_snapshots_managed_total Managed EBS snapshots"
        echo "# TYPE ebs_snapshots_managed_total gauge"
        echo "ebs_snapshots_managed_total{region=\"${AWS_REGION}\"} ${managed_count}"
        echo "# HELP ebs_snapshots_orphaned_total Orphaned EBS snapshots"
        echo "# TYPE ebs_snapshots_orphaned_total gauge"
        echo "ebs_snapshots_orphaned_total{region=\"${AWS_REGION}\"} ${orphan_count}"
        echo "# HELP ebs_snapshots_untagged_total Untagged EBS snapshots"
        echo "# TYPE ebs_snapshots_untagged_total gauge"
        echo "ebs_snapshots_untagged_total{region=\"${AWS_REGION}\"} ${untagged_count}"
        echo "# HELP ebs_snapshots_size_gib_total Total snapshot storage in GiB"
        echo "# TYPE ebs_snapshots_size_gib_total gauge"
        echo "ebs_snapshots_size_gib_total{region=\"${AWS_REGION}\"} ${total_size}"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                         RESTORE MODE
# ══════════════════════════════════════════════════════════════════════

do_restore() {
    if [[ -z "$TARGET_SNAPSHOT" ]]; then
        err "Snapshot ID required. Use --restore SNAP_ID"
        exit 1
    fi

    log "Restoring volume from snapshot ${TARGET_SNAPSHOT}..."

    # Verify snapshot exists and is completed
    local snap_info
    snap_info=$(aws_cmd ec2 describe-snapshots \
        --snapshot-ids "$TARGET_SNAPSHOT" \
        --query 'Snapshots[0].{State:State,Size:VolumeSize,Vol:VolumeId}' \
        --output json 2>/dev/null) || {
        err "Snapshot ${TARGET_SNAPSHOT} not found"
        exit 1
    }

    local snap_state snap_size source_vol
    snap_state=$(echo "$snap_info" | jq -r '.State')
    snap_size=$(echo "$snap_info" | jq -r '.Size')
    source_vol=$(echo "$snap_info" | jq -r '.Vol')

    if [[ "$snap_state" != "completed" ]]; then
        err "Snapshot state is '${snap_state}' — must be 'completed'"
        exit 1
    fi

    # Determine AZ
    if [[ -z "$RESTORE_AZ" ]]; then
        RESTORE_AZ=$(aws_cmd ec2 describe-availability-zones \
            --query 'AvailabilityZones[0].ZoneName' \
            --output text)
        log "No AZ specified, using ${RESTORE_AZ}"
    fi

    # Build create-volume args
    local create_args=(
        ec2 create-volume
        --snapshot-id "$TARGET_SNAPSHOT"
        --availability-zone "$RESTORE_AZ"
        --volume-type "$RESTORE_VOLUME_TYPE"
        --tag-specifications "ResourceType=volume,Tags=[
            {Key=Name,Value=restored-from-${TARGET_SNAPSHOT}},
            {Key=CreatedBy,Value=ebs-snapshot-manager},
            {Key=RestoredFrom,Value=${TARGET_SNAPSHOT}},
            {Key=SourceVolumeId,Value=${source_vol}}
        ]"
    )

    [[ -n "$RESTORE_IOPS" ]] && create_args+=(--iops "$RESTORE_IOPS")
    [[ -n "$RESTORE_THROUGHPUT" ]] && create_args+=(--throughput "$RESTORE_THROUGHPUT")

    local vol_id
    vol_id=$(aws_cmd "${create_args[@]}" \
        --query 'VolumeId' \
        --output text 2>/dev/null) || {
        err "Failed to create volume from snapshot"
        exit 1
    }

    echo -e "  ${GREEN}✓${RESET} Created volume ${vol_id}"
    echo -e "    Source snapshot: ${TARGET_SNAPSHOT}"
    echo -e "    Size:           ${snap_size} GiB"
    echo -e "    Type:           ${RESTORE_VOLUME_TYPE}"
    echo -e "    AZ:             ${RESTORE_AZ}"

    # Wait for volume to become available
    log "Waiting for volume to become available..."
    if aws_cmd ec2 wait volume-available --volume-ids "$vol_id" 2>/dev/null; then
        echo -e "  ${GREEN}✓${RESET} Volume ${vol_id} is available"
    else
        warn "Volume did not become available within timeout"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                          LIST MODE
# ══════════════════════════════════════════════════════════════════════

do_list() {
    local owner_id
    owner_id=$(get_account_id)

    local filters=("Name=owner-id,Values=${owner_id}")

    if [[ -n "$TARGET_VOLUME" ]]; then
        filters+=("Name=volume-id,Values=${TARGET_VOLUME}")
    fi

    local snapshots_json
    snapshots_json=$(aws_cmd ec2 describe-snapshots \
        --owner-ids "$owner_id" \
        ${TARGET_VOLUME:+--filters "Name=volume-id,Values=${TARGET_VOLUME}"} \
        --query 'sort_by(Snapshots, &StartTime) | reverse(@) | [*].{Id:SnapshotId,Vol:VolumeId,Size:VolumeSize,Status:State,Start:StartTime,Desc:Description}' \
        --output json)

    local total
    total=$(echo "$snapshots_json" | jq 'length')

    if [[ "$total" -eq 0 ]]; then
        log "No snapshots found"
        return
    fi

    if [[ "$OUTPUT_FORMAT" == "json" ]]; then
        echo "$snapshots_json" | jq '.'
        return
    fi

    echo ""
    printf "  %-24s %-14s %8s %-12s %-22s %s\n" "SNAPSHOT" "VOLUME" "SIZE" "STATUS" "CREATED" "DESCRIPTION"
    echo "  $(printf '%.0s─' {1..100})"

    echo "$snapshots_json" | jq -c '.[]' | while IFS= read -r snap; do
        local snap_id vol_id size status start_time desc
        snap_id=$(echo "$snap" | jq -r '.Id')
        vol_id=$(echo "$snap" | jq -r '.Vol')
        size=$(echo "$snap" | jq -r '.Size')
        status=$(echo "$snap" | jq -r '.Status')
        start_time=$(echo "$snap" | jq -r '.Start' | cut -c1-19)
        desc=$(echo "$snap" | jq -r '.Desc' | cut -c1-40)

        printf "  %-24s %-14s %6s G %-12s %-22s %s\n" \
            "$snap_id" "$vol_id" "$size" "$status" "$start_time" "$desc"
    done

    echo ""
    log "Total: ${total} snapshot(s)"
}

# ══════════════════════════════════════════════════════════════════════
#                          MAIN
# ══════════════════════════════════════════════════════════════════════

show_help() {
    cat <<EOF
Usage: $SCRIPT_NAME [MODE] [OPTIONS]

Manage AWS EBS snapshots — create, prune, copy, audit, restore, and list.

MODES:
    --snapshot              Create snapshots of EBS volumes
    --prune                 Delete snapshots older than retention period
    --copy-region REGION    Copy latest snapshots to another region
    --audit                 Snapshot health check and cost report
    --restore SNAP_ID       Create a new volume from a snapshot
    --list                  List snapshots

OPTIONS:
    --volume-id VOL_ID      Target a specific volume
    --tag KEY=VALUE         Filter volumes by tag (e.g., --tag Backup=true)
    --retention-days N      Retention period in days (default: $RETENTION_DAYS)
    --format FORMAT         Output: text (default), json, prometheus
    --force                 Actually delete in prune mode (default is dry-run)
    --no-wait               Don't wait for snapshot completion
    --verbose               Show debug output
    --no-color              Disable colored output
    --help, -h              Show this help

ENVIRONMENT VARIABLES:
    AWS_PROFILE             AWS CLI profile to use
    AWS_REGION              AWS region (or set via aws configure)
    VOLUME_IDS              Comma-separated volume IDs to snapshot
    VOLUME_TAG_KEY          Filter volumes by tag key
    VOLUME_TAG_VALUE        Filter volumes by tag value
    RETENTION_DAYS          Days to retain snapshots (default: 30)
    COPY_TO_REGION          Target region for cross-region copy
    SNAPSHOT_DESCRIPTION    Description for new snapshots
    NO_WAIT                 Don't wait for snapshot completion (true/false)
    DRY_RUN                 Dry-run for prune mode (default: true)
    RESTORE_AZ              Availability zone for restored volume
    RESTORE_VOLUME_TYPE     Volume type for restore (default: gp3)
    RESTORE_IOPS            IOPS for restored volume (gp3/io1/io2)
    RESTORE_THROUGHPUT      Throughput for restored volume (gp3)
    OUTPUT_FORMAT           Output format: text, json, prometheus
    VERBOSE                 Debug output (true/false)
    COLOR                   Color: auto, always, never

EXAMPLES:
    # Snapshot all volumes in the region
    ./$(basename "$0") --snapshot

    # Snapshot specific volume
    ./$(basename "$0") --snapshot --volume-id vol-0123456789abcdef0

    # Snapshot volumes tagged Backup=true
    ./$(basename "$0") --snapshot --tag Backup=true

    # Dry-run prune (see what would be deleted)
    ./$(basename "$0") --prune --retention-days 14

    # Actually delete old snapshots
    ./$(basename "$0") --prune --retention-days 14 --force

    # Copy latest snapshots to us-west-2
    ./$(basename "$0") --copy-region us-west-2

    # Audit snapshots (orphans, cost, untagged)
    ./$(basename "$0") --audit

    # Audit with Prometheus metrics output
    ./$(basename "$0") --audit --format prometheus

    # Restore volume from snapshot
    ./$(basename "$0") --restore snap-0123456789abcdef0

    # Restore with specific AZ and volume type
    RESTORE_AZ=us-east-1a RESTORE_VOLUME_TYPE=io2 RESTORE_IOPS=5000 \\
      ./$(basename "$0") --restore snap-0123456789abcdef0

    # List all snapshots
    ./$(basename "$0") --list

    # List snapshots for a specific volume
    ./$(basename "$0") --list --volume-id vol-0123456789abcdef0
EOF
}

main() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --snapshot)         RUN_MODE="snapshot" ;;
            --prune)            RUN_MODE="prune" ;;
            --copy-region)      RUN_MODE="copy-region"; COPY_TO_REGION="$2"; shift ;;
            --audit)            RUN_MODE="audit" ;;
            --restore)          RUN_MODE="restore"; TARGET_SNAPSHOT="$2"; shift ;;
            --list)             RUN_MODE="list" ;;
            --volume-id)        TARGET_VOLUME="$2"; VOLUME_IDS="$2"; shift ;;
            --tag)
                local tag_pair="$2"
                VOLUME_TAG_KEY="${tag_pair%%=*}"
                VOLUME_TAG_VALUE="${tag_pair#*=}"
                shift
                ;;
            --retention-days)   RETENTION_DAYS="$2"; shift ;;
            --format)           OUTPUT_FORMAT="$2"; shift ;;
            --force)            DRY_RUN="false" ;;
            --no-wait)          NO_WAIT="true" ;;
            --verbose)          VERBOSE="true" ;;
            --no-color)         COLOR="never" ;;
            --help|-h)          show_help; exit 0 ;;
            *)                  err "Unknown option: $1"; echo ""; show_help; exit 1 ;;
        esac
        shift
    done

    setup_colors

    if [[ -z "$RUN_MODE" ]]; then
        err "No mode specified. Use --snapshot, --prune, --audit, --list, --restore, or --copy-region"
        echo ""
        show_help
        exit 1
    fi

    START_TIME=$(date +%s)

    echo ""
    echo -e "${BOLD}EBS Snapshot Manager${RESET}"
    echo -e "Region: ${AWS_REGION:-$(aws configure get region 2>/dev/null || echo 'default')}"
    echo -e "Mode:   ${RUN_MODE}"
    echo -e "Time:   $(date -u +%Y-%m-%dT%H:%M:%SZ)"
    echo ""

    check_deps

    case "$RUN_MODE" in
        snapshot)     do_snapshot ;;
        prune)        do_prune ;;
        copy-region)  do_copy_region ;;
        audit)        do_audit ;;
        restore)      do_restore ;;
        list)         do_list ;;
    esac

    local end_time
    end_time=$(date +%s)
    local duration=$(( end_time - START_TIME ))
    echo ""
    log "Completed in ${duration}s"

    if [[ $WARNINGS -gt 0 ]]; then
        exit 2
    fi
}

main "$@"