#!/bin/bash

###############################################################################
# Glacier Deep Archive Upload Script — Resumable Multipart Upload
#
# Uploads pre-built archives from transit disks to AWS S3 Glacier Deep Archive
# using the S3 multipart upload API. Each 100MB part is tracked individually,
# so a power outage loses at most ~100MB of upload progress.
#
# Two levels of resume:
#   1. Whole-archive: completed.txt tracks fully uploaded archives
#   2. Per-part: .upload state files track individual 100MB parts
#
# Usage:
#   ./glacier_upload.sh [--dry-run] [--yes] [--cleanup]
#
# Options:
#   --dry-run   Show upload plan, no API calls
#   --yes       Skip confirmation prompt (for background/nohup execution)
#   --cleanup   Abort incomplete S3 multipart uploads, clean local state files
###############################################################################

set -euo pipefail

# Configuration - EDIT THESE VALUES
S3_BUCKET="your-bucket-name"    # Must be changed before use
S3_PREFIX="photo-video-archive"
STORAGE_CLASS="DEEP_ARCHIVE"
TRANSIT_DISK_1="/mnt/sdd_transit"
TRANSIT_DISK_2="/mnt/sdc_transit"
PART_SIZE=$((256 * 1024 * 1024))  # 100MB per part (~1024 parts for 100GB, under 10000 limit)
UPLOAD_STATE_DIR_NAME=".upload_state"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m'

# Flags
DRY_RUN=false
AUTO_YES=false
CLEANUP=false
for arg in "$@"; do
    case "$arg" in
        --dry-run) DRY_RUN=true ;;
        --yes)     AUTO_YES=true ;;
        --cleanup) CLEANUP=true ;;
        *) echo "ERROR: Unknown option: $arg" >&2; exit 1 ;;
    esac
done

# Log directory (on first available transit disk)
LOG_DIR=""
LOG_FILE=""

###############################################################################
# Utility Functions
###############################################################################

log() {
    local msg="[$(date '+%Y-%m-%d %H:%M:%S')] $1"
    echo "$msg"
    [ -n "$LOG_FILE" ] && echo "$msg" >> "$LOG_FILE" || true
}

log_color() {
    local msg="[$(date '+%Y-%m-%d %H:%M:%S')] $1"
    echo -e "${2}${msg}${NC}"
    [ -n "$LOG_FILE" ] && echo "$msg" >> "$LOG_FILE" || true
}

error_exit() {
    log_color "ERROR: $1" "$RED" >&2
    exit 1
}

human_size() {
    local bytes=$1
    if (( bytes >= 1099511627776 )); then
        echo "$(echo "scale=1; $bytes / 1099511627776" | bc)T"
    elif (( bytes >= 1073741824 )); then
        echo "$(echo "scale=1; $bytes / 1073741824" | bc)G"
    elif (( bytes >= 1048576 )); then
        echo "$(echo "scale=1; $bytes / 1048576" | bc)M"
    else
        echo "${bytes}B"
    fi
}

setup_logging() {
    for disk in "$TRANSIT_DISK_1" "$TRANSIT_DISK_2"; do
        if [ -d "$disk" ]; then
            LOG_DIR="$disk/$UPLOAD_STATE_DIR_NAME"
            mkdir -p "$LOG_DIR"
            LOG_FILE="$LOG_DIR/upload_$(date +%Y%m%d_%H%M%S).log"
            return
        fi
    done
    error_exit "No transit disk is accessible"
}

###############################################################################
# Scanning & Planning
###############################################################################

scan_archives() {
    # Find all .tar.gz files on both transit disks, output: disk_path\tarchive_name\tsize_bytes
    # Sorted by archive name (natural sort) for consistent ordering
    local tmp
    tmp=$(mktemp)

    for disk in "$TRANSIT_DISK_1" "$TRANSIT_DISK_2"; do
        [ ! -d "$disk" ] && continue
        while IFS= read -r -d '' archive_path; do
            local name size_bytes
            name=$(basename "$archive_path")
            size_bytes=$(stat -c%s "$archive_path" 2>/dev/null || stat -f%z "$archive_path" 2>/dev/null)
            echo -e "${disk}\t${name}\t${size_bytes}" >> "$tmp"
        done < <(find "$disk" -maxdepth 1 -name "*.tar.gz" -print0)
    done

    # Natural sort by archive name
    sort -t$'\t' -k2,2V "$tmp"
    rm -f "$tmp"
}

is_upload_completed() {
    local disk_path="$1"
    local archive_name="$2"
    local completed_file="$disk_path/$UPLOAD_STATE_DIR_NAME/completed.txt"
    grep -Fxq "$archive_name" "$completed_file" 2>/dev/null
}

mark_upload_completed() {
    local disk_path="$1"
    local archive_name="$2"
    local completed_file="$disk_path/$UPLOAD_STATE_DIR_NAME/completed.txt"
    echo "$archive_name" >> "$completed_file"
}

display_plan() {
    # Input: lines of disk_path\tarchive_name\tsize_bytes from scan_archives
    local plan_data="$1"

    local total_count=0 total_size=0
    local completed_count=0 completed_size=0
    local pending_count=0 pending_size=0
    local resumable_count=0

    echo ""
    log_color "==================== UPLOAD PLAN ====================" "$CYAN"
    echo ""
    echo -e "  ${BLUE}S3 destination:${NC} s3://$S3_BUCKET/$S3_PREFIX/"
    echo -e "  ${BLUE}Storage class:${NC}  $STORAGE_CLASS"
    echo -e "  ${BLUE}Part size:${NC}      $(human_size $PART_SIZE)"
    echo ""

    for disk in "$TRANSIT_DISK_1" "$TRANSIT_DISK_2"; do
        [ ! -d "$disk" ] && continue
        local disk_name
        disk_name=$(basename "$disk")
        local disk_archives=0 disk_size=0 disk_done=0 disk_resume=0 disk_pending=0

        echo -e "${BLUE}${disk_name} (${disk}):${NC}"

        while IFS=$'\t' read -r d_path a_name a_size; do
            [ "$d_path" != "$disk" ] && continue
            disk_archives=$((disk_archives + 1))
            disk_size=$((disk_size + a_size))
            total_count=$((total_count + 1))
            total_size=$((total_size + a_size))

            local state_file="$disk/$UPLOAD_STATE_DIR_NAME/${a_name}.upload"
            local total_parts=$(( (a_size + PART_SIZE - 1) / PART_SIZE ))

            if is_upload_completed "$disk" "$a_name"; then
                echo -e "  ${GREEN}done${NC}    $(human_size "$a_size")\t$a_name"
                completed_count=$((completed_count + 1))
                completed_size=$((completed_size + a_size))
                disk_done=$((disk_done + 1))
            elif [ -f "$state_file" ]; then
                local parts_done
                parts_done=$(grep -c $'^\t*[0-9]' "$state_file" 2>/dev/null || echo 0)
                echo -e "  ${YELLOW}resume${NC}  $(human_size "$a_size")\t$a_name (${parts_done}/${total_parts} parts)"
                resumable_count=$((resumable_count + 1))
                pending_count=$((pending_count + 1))
                pending_size=$((pending_size + a_size))
                disk_resume=$((disk_resume + 1))
            else
                echo -e "  ${CYAN}new${NC}     $(human_size "$a_size")\t$a_name (${total_parts} parts)"
                pending_count=$((pending_count + 1))
                pending_size=$((pending_size + a_size))
                disk_pending=$((disk_pending + 1))
            fi
        done <<< "$plan_data"

        if [ "$disk_archives" -eq 0 ]; then
            echo "  (no archives)"
        fi
        echo -e "  ${YELLOW}--- $disk_archives archives, $(human_size $disk_size) total ---${NC}"
        echo ""
    done

    echo -e "${CYAN}========================================================${NC}"
    echo -e "${CYAN}Total: $total_count archives, $(human_size $total_size)${NC}"
    echo -e "${GREEN}  Completed: $completed_count ($(human_size $completed_size))${NC}"
    if [ "$resumable_count" -gt 0 ]; then
        echo -e "${YELLOW}  Resumable: $resumable_count${NC}"
    fi
    echo -e "${CYAN}  Pending:   $pending_count ($(human_size $pending_size))${NC}"
    echo -e "${CYAN}========================================================${NC}"
    echo ""
}

###############################################################################
# Multipart Upload
###############################################################################

upload_archive() {
    local disk_path="$1"
    local archive_name="$2"
    local file_size="$3"

    local archive_path="$disk_path/$archive_name"
    local s3_key="$S3_PREFIX/$archive_name"
    local state_dir="$disk_path/$UPLOAD_STATE_DIR_NAME"
    local state_file="$state_dir/${archive_name}.upload"
    local disk_name
    disk_name=$(basename "$disk_path")

    mkdir -p "$state_dir"

    # Skip if already completed
    if is_upload_completed "$disk_path" "$archive_name"; then
        log_color "SKIP: $archive_name (already completed)" "$YELLOW"
        return 0
    fi

    local total_parts=$(( (file_size + PART_SIZE - 1) / PART_SIZE ))

    echo ""
    log_color "================================================================" "$BLUE"
    log_color "$archive_name from $disk_name ($(human_size "$file_size"), $total_parts parts)" "$BLUE"
    log_color "================================================================" "$BLUE"

    # Step 1: Initiate or resume multipart upload
    local upload_id=""

    if [ -f "$state_file" ]; then
        # Read existing upload_id from state file
        upload_id=$(grep '^upload_id=' "$state_file" | head -1 | cut -d= -f2-)

        if [ -n "$upload_id" ]; then
            # Validate the upload is still alive on S3
            log_color "Resuming upload: $upload_id" "$YELLOW"
            if ! aws s3api list-parts \
                    --bucket "$S3_BUCKET" \
                    --key "$s3_key" \
                    --upload-id "$upload_id" \
                    --max-parts 0 \
                    --query 'UploadId' \
                    --output text &>/dev/null; then
                log_color "Previous upload expired on S3, starting fresh" "$YELLOW"
                upload_id=""
                rm -f "$state_file"
            fi
        fi
    fi

    if [ -z "$upload_id" ]; then
        log_color "Initiating new multipart upload..." "$BLUE"
        upload_id=$(aws s3api create-multipart-upload \
            --bucket "$S3_BUCKET" \
            --key "$s3_key" \
            --storage-class "$STORAGE_CLASS" \
            --query 'UploadId' \
            --output text)

        if [ -z "$upload_id" ] || [ "$upload_id" = "None" ]; then
            log_color "FAIL: Could not initiate multipart upload for $archive_name" "$RED"
            return 1
        fi

        # Write state file header
        {
            echo "upload_id=$upload_id"
            echo "s3_key=$s3_key"
            echo "file_size=$file_size"
            echo "part_size=$PART_SIZE"
            echo "total_parts=$total_parts"
        } > "$state_file"

        log_color "Upload initiated: $upload_id" "$GREEN"
    fi

    # Step 2: Upload parts
    # Determine which parts are already done from state file
    local -A completed_parts=()
    while IFS=$'\t' read -r pnum petag; do
        [ -z "$pnum" ] && continue
        # Only consider lines that start with a number (part records)
        if [[ "$pnum" =~ ^[0-9]+$ ]]; then
            completed_parts[$pnum]="$petag"
        fi
    done < "$state_file"

    local parts_uploaded=${#completed_parts[@]}
    local upload_start_time
    upload_start_time=$(date +%s)
    local bytes_uploaded_session=0

    for (( part_num=1; part_num<=total_parts; part_num++ )); do
        # Skip already completed parts
        if [ -n "${completed_parts[$part_num]:-}" ]; then
            continue
        fi

        # Calculate offset and this part's size
        local offset=$(( (part_num - 1) * PART_SIZE ))
        local this_part_size=$PART_SIZE
        if (( offset + this_part_size > file_size )); then
            this_part_size=$(( file_size - offset ))
        fi

        # Extract part to temp file, then upload (100MB temp space)
        local part_file="$state_dir/.part_tmp"
        dd if="$archive_path" bs=$PART_SIZE skip=$((part_num - 1)) count=1 iflag=fullblock 2>/dev/null > "$part_file"

        local etag
        etag=$(aws s3api upload-part \
            --bucket "$S3_BUCKET" \
            --key "$s3_key" \
            --upload-id "$upload_id" \
            --part-number "$part_num" \
            --body "$part_file" \
            --query 'ETag' \
            --output text 2>/dev/null)

        rm -f "$part_file"

        if [ -z "$etag" ] || [ "$etag" = "None" ]; then
            log_color "FAIL: Upload failed for part $part_num of $archive_name" "$RED"
            return 1
        fi

        # Append part record to state file (immediate durability)
        printf '%s\t%s\n' "$part_num" "$etag" >> "$state_file"
        parts_uploaded=$((parts_uploaded + 1))
        bytes_uploaded_session=$((bytes_uploaded_session + this_part_size))

        # Inline progress
        local pct=$(( parts_uploaded * 100 / total_parts ))
        local elapsed=$(( $(date +%s) - upload_start_time ))
        local speed_str=""
        if (( elapsed > 0 && bytes_uploaded_session > 0 )); then
            local speed_mbps=$(( bytes_uploaded_session / elapsed / 1048576 ))
            speed_str=" ${speed_mbps} MB/s"
        fi
        printf "\r  Part %d/%d (%d%%)%s    " "$parts_uploaded" "$total_parts" "$pct" "$speed_str"
    done
    echo ""

    # Step 3: Complete multipart upload
    log_color "Completing multipart upload..." "$BLUE"

    # Build the JSON for complete-multipart-upload from state file
    # Format: {"Parts":[{"PartNumber":1,"ETag":"\"abc\""},...]}"
    local parts_json
    parts_json=$(awk -F'\t' '
        /^[0-9]+\t/ {
            if (n++) printf ","
            printf "{\"PartNumber\":%s,\"ETag\":%s}", $1, $2
        }
        BEGIN { printf "{\"Parts\":[" }
        END { printf "]}" }
    ' "$state_file")

    local complete_result
    complete_result=$(aws s3api complete-multipart-upload \
        --bucket "$S3_BUCKET" \
        --key "$s3_key" \
        --upload-id "$upload_id" \
        --multipart-upload "$parts_json" \
        --query 'Location' \
        --output text 2>&1) || {
        log_color "FAIL: complete-multipart-upload failed for $archive_name" "$RED"
        log "$complete_result"
        return 1
    }

    log_color "Multipart upload completed" "$GREEN"

    # Step 4: Verify
    log_color "Verifying upload..." "$BLUE"
    local s3_size
    s3_size=$(aws s3api head-object \
        --bucket "$S3_BUCKET" \
        --key "$s3_key" \
        --query 'ContentLength' \
        --output text 2>/dev/null)

    if [ -z "$s3_size" ] || [ "$s3_size" = "None" ]; then
        log_color "FAIL: Verification failed — object not found on S3" "$RED"
        return 1
    fi

    if [ "$s3_size" -ne "$file_size" ]; then
        log_color "FAIL: Size mismatch! Local: $file_size, S3: $s3_size" "$RED"
        return 1
    fi

    log_color "Verified: size matches ($file_size bytes)" "$GREEN"

    # Step 5: Clean up
    mark_upload_completed "$disk_path" "$archive_name"
    rm -f "$state_file"

    log_color "Deleting local archive..." "$BLUE"
    rm -f "$archive_path"
    log_color "DONE: $archive_name" "$GREEN"

    return 0
}

upload_master_manifests() {
    log_color "Uploading master manifests to S3..." "$BLUE"

    for disk in "$TRANSIT_DISK_1" "$TRANSIT_DISK_2"; do
        [ ! -d "$disk" ] && continue
        local manifest="$disk/archive_master_manifest.txt"
        [ ! -f "$manifest" ] && continue

        local disk_name
        disk_name=$(basename "$disk")
        local s3_key="$S3_PREFIX/${disk_name}_master_manifest.txt"

        if aws s3 cp "$manifest" "s3://$S3_BUCKET/$s3_key" \
            --content-type "text/plain" \
            2>&1 | tee -a "${LOG_FILE:-/dev/null}"; then
            log_color "Uploaded: $disk_name master manifest -> s3://$S3_BUCKET/$s3_key" "$GREEN"
        else
            log_color "Warning: Failed to upload manifest from $disk_name" "$YELLOW"
        fi
    done
}

###############################################################################
# Cleanup
###############################################################################

cleanup_multipart_uploads() {
    log_color "========================================" "$BLUE"
    log_color "Cleaning up multipart uploads" "$BLUE"
    log_color "========================================" "$BLUE"
    echo ""

    # List and abort incomplete multipart uploads on S3
    log_color "Checking S3 for incomplete multipart uploads..." "$BLUE"
    local uploads
    uploads=$(aws s3api list-multipart-uploads \
        --bucket "$S3_BUCKET" \
        --prefix "$S3_PREFIX/" \
        --query 'Uploads[].{Key:Key,UploadId:UploadId,Initiated:Initiated}' \
        --output text 2>/dev/null)

    if [ -n "$uploads" ] && [ "$uploads" != "None" ]; then
        local upload_count=0
        while IFS=$'\t' read -r initiated key upload_id; do
            [ -z "$key" ] && continue
            echo -e "  ${YELLOW}$key${NC} (started: $initiated)"
            echo "    Upload ID: $upload_id"
            upload_count=$((upload_count + 1))
        done <<< "$uploads"

        if [ "$upload_count" -gt 0 ]; then
            echo ""
            if ! $AUTO_YES; then
                read -p "Abort all $upload_count incomplete uploads? (yes/no): " confirm
                if [[ "$confirm" != "yes" ]]; then
                    log "Skipped S3 cleanup"
                    echo ""
                else
                    while IFS=$'\t' read -r initiated key upload_id; do
                        [ -z "$key" ] && continue
                        if aws s3api abort-multipart-upload \
                                --bucket "$S3_BUCKET" \
                                --key "$key" \
                                --upload-id "$upload_id" 2>/dev/null; then
                            log_color "Aborted: $key" "$GREEN"
                        else
                            log_color "Failed to abort: $key" "$RED"
                        fi
                    done <<< "$uploads"
                fi
            else
                while IFS=$'\t' read -r initiated key upload_id; do
                    [ -z "$key" ] && continue
                    if aws s3api abort-multipart-upload \
                            --bucket "$S3_BUCKET" \
                            --key "$key" \
                            --upload-id "$upload_id" 2>/dev/null; then
                        log_color "Aborted: $key" "$GREEN"
                    else
                        log_color "Failed to abort: $key" "$RED"
                    fi
                done <<< "$uploads"
            fi
        fi
    else
        log_color "No incomplete multipart uploads found on S3" "$GREEN"
    fi

    echo ""

    # Clean local state files
    log_color "Cleaning local upload state..." "$BLUE"
    for disk in "$TRANSIT_DISK_1" "$TRANSIT_DISK_2"; do
        [ ! -d "$disk" ] && continue
        local disk_name state_dir
        disk_name=$(basename "$disk")
        state_dir="$disk/$UPLOAD_STATE_DIR_NAME"

        if [ ! -d "$state_dir" ]; then
            log_color "  $disk_name: no state directory" "$GREEN"
            continue
        fi

        local upload_files
        upload_files=$(find "$state_dir" -name "*.upload" 2>/dev/null | wc -l | tr -d ' ')
        local has_completed=false
        [ -f "$state_dir/completed.txt" ] && has_completed=true

        echo -e "  ${BLUE}$disk_name:${NC}"
        echo "    Upload state files: $upload_files"
        echo "    Completed tracking: $has_completed"

        if (( upload_files == 0 )) && ! $has_completed; then
            log_color "    Nothing to clean" "$GREEN"
            continue
        fi

        if ! $AUTO_YES; then
            read -p "    Delete state files from $disk_name? (yes/no): " confirm
            if [[ "$confirm" != "yes" ]]; then
                log "    Skipped $disk_name"
                continue
            fi
        fi

        find "$state_dir" -name "*.upload" -delete 2>/dev/null
        rm -f "$state_dir/completed.txt"
        # Remove log files too
        find "$state_dir" -name "upload_*.log" -delete 2>/dev/null
        rmdir "$state_dir" 2>/dev/null || true
        log_color "    Cleaned $disk_name" "$GREEN"
    done

    echo ""
    log_color "Cleanup complete" "$GREEN"
}

###############################################################################
# Prerequisites
###############################################################################

check_prerequisites() {
    log_color "Checking prerequisites..." "$BLUE"

    if ! command -v aws &>/dev/null; then
        error_exit "AWS CLI is not installed. Install it with: brew install awscli (macOS) or see AWS docs (Linux)"
    fi

    if ! aws sts get-caller-identity &>/dev/null; then
        error_exit "AWS credentials not configured. Run: aws configure"
    fi

    if ! aws s3 ls "s3://$S3_BUCKET" &>/dev/null; then
        error_exit "S3 bucket '$S3_BUCKET' does not exist or is not accessible"
    fi

    # Check at least one transit disk exists
    local disk_found=false
    for disk in "$TRANSIT_DISK_1" "$TRANSIT_DISK_2"; do
        [ -d "$disk" ] && disk_found=true
    done
    if ! $disk_found; then
        error_exit "No transit disk found. Check TRANSIT_DISK_1 and TRANSIT_DISK_2 paths."
    fi

    log_color "All prerequisites met!" "$GREEN"
}

###############################################################################
# Main
###############################################################################

main() {
    # Handle --cleanup before anything else
    if $CLEANUP; then
        setup_logging
        check_prerequisites
        cleanup_multipart_uploads
        exit 0
    fi

    log_color "========================================" "$BLUE"
    log_color "Glacier Deep Archive Upload — Resumable Multipart" "$BLUE"
    log_color "========================================" "$BLUE"
    echo ""

    check_prerequisites
    setup_logging

    # Scan archives on transit disks
    log_color "Scanning transit disks for archives..." "$BLUE"
    local plan_data
    plan_data=$(scan_archives)

    if [ -z "$plan_data" ]; then
        log_color "No archives found on transit disks" "$YELLOW"
        echo "Expected .tar.gz files in: $TRANSIT_DISK_1 and/or $TRANSIT_DISK_2"
        exit 0
    fi

    display_plan "$plan_data"

    if $DRY_RUN; then
        log_color "DRY RUN mode — no uploads will be performed" "$YELLOW"
        exit 0
    fi

    # Confirmation
    echo -e "${YELLOW}This script will:${NC}"
    echo "  1. Upload archives to: s3://$S3_BUCKET/$S3_PREFIX/"
    echo "  2. Storage class: $STORAGE_CLASS"
    echo "  3. Part size: $(human_size $PART_SIZE) (resumable per-part)"
    echo "  4. Delete local archive after successful upload + verification"
    echo "  5. Upload master manifests at the end"
    echo ""
    echo -e "${RED}WARNING: Make sure your S3 bucket name is correct!${NC}"
    echo -e "${RED}Glacier Deep Archive has retrieval costs and delays.${NC}"
    echo ""

    if ! $AUTO_YES; then
        read -p "Continue? (yes/no): " confirm
        if [[ "$confirm" != "yes" ]]; then
            log "Operation cancelled by user"
            exit 0
        fi
    else
        log "Auto-confirmed (--yes)"
    fi

    log_color "Starting uploads..." "$GREEN"
    echo ""

    # Initialize state dirs
    for disk in "$TRANSIT_DISK_1" "$TRANSIT_DISK_2"; do
        [ ! -d "$disk" ] && continue
        mkdir -p "$disk/$UPLOAD_STATE_DIR_NAME"
        touch "$disk/$UPLOAD_STATE_DIR_NAME/completed.txt"
    done

    # Process each archive
    local success_count=0
    local error_count=0
    local skip_count=0

    while IFS=$'\t' read -r disk_path archive_name file_size; do
        if is_upload_completed "$disk_path" "$archive_name"; then
            skip_count=$((skip_count + 1))
            continue
        fi
        if upload_archive "$disk_path" "$archive_name" "$file_size"; then
            success_count=$((success_count + 1))
        else
            error_count=$((error_count + 1))
            log_color "ERROR: Failed to upload $archive_name" "$RED"
        fi
    done <<< "$plan_data"

    # Upload master manifests
    echo ""
    upload_master_manifests

    # Final summary
    echo ""
    log_color "========================================" "$BLUE"
    log_color "UPLOAD COMPLETE" "$BLUE"
    log_color "========================================" "$BLUE"
    log "Uploaded:  $success_count"
    log "Skipped:   $skip_count (already done)"
    log "Failed:    $error_count"
    log "S3 Bucket: s3://$S3_BUCKET/$S3_PREFIX/"
    [ -n "$LOG_FILE" ] && log "Log file:  $LOG_FILE"
    log_color "========================================" "$BLUE"

    return $error_count
}

main "$@"
