#!/usr/bin/env bash

#########################################################################################
#### aws-cost-reporter.sh — Daily AWS cost breakdown by service, account, or tag     ####
#### Supports email (SES), Slack webhooks, CSV/JSON export, period comparison        ####
#### Requires: bash 4+, aws-cli v2, jq                                              ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.00                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   export AWS_PROFILE="billing"                                                  ####
####   ./aws-cost-reporter.sh --daily                                                ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
AWS_REGION="${AWS_REGION:-us-east-1}"
GROUP_BY="${GROUP_BY:-SERVICE}"
OUTPUT_FORMAT="${OUTPUT_FORMAT:-text}"
SES_FROM_ADDRESS="${SES_FROM_ADDRESS:-}"
SLACK_WEBHOOK_URL="${SLACK_WEBHOOK_URL:-}"
COST_TAG_KEY="${COST_TAG_KEY:-}"
COST_TAG_VALUE="${COST_TAG_VALUE:-}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
RUN_MODE=""
CUSTOM_START=""
CUSTOM_END=""
EMAIL_TO=""
SLACK_URL=""
START_TIME=""

# ── Colors ────────────────────────────────────────────────────────────
RED="" GREEN="" YELLOW="" BLUE="" BOLD="" DIM="" RESET=""
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" DIM="" RESET=""
        return
    fi
    if [[ "$COLOR" == "auto" && ! -t 1 ]]; then
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" DIM="" RESET=""
        return
    fi
    RED="\033[0;31m"
    GREEN="\033[0;32m"
    YELLOW="\033[0;33m"
    # shellcheck disable=SC2034 # BLUE reserved for future use / caller scripts
    BLUE="\033[0;34m"
    # shellcheck disable=SC2034 # BOLD reserved for future use / caller scripts
    BOLD="\033[1m"
    DIM="\033[2m"
    RESET="\033[0m"
}

# ── Logging ───────────────────────────────────────────────────────────
log_info()  { printf "${GREEN}[INFO]${RESET} %s\n" "$*"; }
log_warn()  { printf "${YELLOW}[WARN]${RESET} %s\n" "$*" >&2; }
log_error() { printf "${RED}[ERROR]${RESET} %s\n" "$*" >&2; }
log_debug() { [[ "$VERBOSE" == "true" ]] && printf "${DIM}[DEBUG] %s${RESET}\n" "$*"; }

# ── Helpers ───────────────────────────────────────────────────────────
die() { log_error "$@"; exit 1; }

check_deps() {
    local missing=()
    command -v aws  >/dev/null 2>&1 || missing+=("aws-cli")
    command -v jq   >/dev/null 2>&1 || missing+=("jq")
    command -v curl >/dev/null 2>&1 || missing+=("curl")
    if (( ${#missing[@]} > 0 )); then
        die "Missing required tools: ${missing[*]}"
    fi

    local bash_major="${BASH_VERSINFO[0]}"
    if (( bash_major < 4 )); then
        die "Requires bash 4+, found ${BASH_VERSION}"
    fi
}

validate_date() {
    local d="$1"
    if [[ ! "$d" =~ ^[0-9]{4}-[0-9]{2}-[0-9]{2}$ ]]; then
        die "Invalid date format: $d (expected YYYY-MM-DD)"
    fi
}

# ── Date math (portable) ─────────────────────────────────────────────
date_offset() {
    # Usage: date_offset YYYY-MM-DD -N  → date N days before
    local base="$1" offset="$2"
    if date --version >/dev/null 2>&1; then
        # GNU date
        date -d "${base} ${offset} days" +%Y-%m-%d
    else
        # macOS date
        date -j -v"${offset}d" -f "%Y-%m-%d" "$base" +%Y-%m-%d
    fi
}

today_utc() { date -u +%Y-%m-%d; }

first_of_month() {
    local d="$1"
    echo "${d:0:8}01"
}

first_of_prev_month() {
    local d="$1"
    local year="${d:0:4}"
    local month="${d:5:2}"
    month=$((10#$month - 1))
    if (( month == 0 )); then
        month=12
        year=$((year - 1))
    fi
    printf "%04d-%02d-01" "$year" "$month"
}

days_between() {
    local s="$1" e="$2"
    local ss se
    if date --version >/dev/null 2>&1; then
        ss=$(date -d "$s" +%s)
        se=$(date -d "$e" +%s)
    else
        ss=$(date -j -f "%Y-%m-%d" "$s" +%s)
        se=$(date -j -f "%Y-%m-%d" "$e" +%s)
    fi
    echo $(( (se - ss) / 86400 ))
}

# ── Compute date ranges ──────────────────────────────────────────────
compute_ranges() {
    local today
    today="$(today_utc)"

    case "$RUN_MODE" in
        daily)
            PERIOD_START="$(date_offset "$today" -1)"
            PERIOD_END="$today"
            PREV_START="$(date_offset "$today" -2)"
            PREV_END="$(date_offset "$today" -1)"
            ;;
        weekly)
            PERIOD_START="$(date_offset "$today" -7)"
            PERIOD_END="$today"
            PREV_START="$(date_offset "$today" -14)"
            PREV_END="$(date_offset "$today" -7)"
            ;;
        monthly)
            PERIOD_START="$(first_of_month "$today")"
            PERIOD_END="$today"
            local prev_first
            prev_first="$(first_of_prev_month "$today")"
            PREV_START="$prev_first"
            PREV_END="$PERIOD_START"
            ;;
        custom)
            PERIOD_START="$CUSTOM_START"
            PERIOD_END="$CUSTOM_END"
            local span
            span="$(days_between "$CUSTOM_START" "$CUSTOM_END")"
            PREV_START="$(date_offset "$CUSTOM_START" "-$span")"
            PREV_END="$CUSTOM_START"
            ;;
        *)
            die "Unknown mode: $RUN_MODE"
            ;;
    esac

    log_debug "Current period: $PERIOD_START → $PERIOD_END"
    log_debug "Previous period: $PREV_START → $PREV_END"
}

# ── Build Cost Explorer request ───────────────────────────────────────
build_ce_filter() {
    local filter=""
    if [[ -n "$COST_TAG_KEY" && -n "$COST_TAG_VALUE" ]]; then
        filter=$(cat <<EOF
{
    "Tags": {
        "Key": "$COST_TAG_KEY",
        "Values": ["$COST_TAG_VALUE"],
        "MatchOptions": ["EQUALS"]
    }
}
EOF
        )
    fi
    echo "$filter"
}

build_group_by() {
    case "$GROUP_BY" in
        SERVICE)
            echo '[{"Type":"DIMENSION","Key":"SERVICE"}]'
            ;;
        TAG)
            if [[ -z "$COST_TAG_KEY" ]]; then
                die "--group-by TAG requires --tag KEY=VALUE"
            fi
            echo "[{\"Type\":\"TAG\",\"Key\":\"$COST_TAG_KEY\"}]"
            ;;
        LINKED_ACCOUNT)
            echo '[{"Type":"DIMENSION","Key":"LINKED_ACCOUNT"}]'
            ;;
        *)
            die "Invalid --group-by value: $GROUP_BY (expected SERVICE, TAG, or LINKED_ACCOUNT)"
            ;;
    esac
}

# ── Query Cost Explorer ───────────────────────────────────────────────
query_costs() {
    local start="$1" end="$2"
    local group_by

    group_by="$(build_group_by)"
    local filter
    filter="$(build_ce_filter)"

    local cmd=(
        aws ce get-cost-and-usage
        --region "$AWS_REGION"
        --time-period "Start=${start},End=${end}"
        --granularity MONTHLY
        --metrics BlendedCost
        --group-by "$group_by"
        --output json
    )

    if [[ -n "$filter" ]]; then
        cmd+=(--filter "$filter")
    fi

    log_debug "Running: ${cmd[*]}"
    "${cmd[@]}" 2>/dev/null
}

# ── Parse cost data ──────────────────────────────────────────────────
parse_costs() {
    local raw="$1"
    echo "$raw" | jq -r '
        [.ResultsByTime[].Groups[] |
            {
                key: .Keys[0],
                amount: (.Metrics.BlendedCost.Amount | tonumber)
            }
        ] |
        group_by(.key) |
        map({
            key: .[0].key,
            total: (map(.amount) | add)
        }) |
        sort_by(-.total) |
        .[] |
        "\(.key)\t\(.total)"
    ' 2>/dev/null || echo ""
}

# ── Format helpers ────────────────────────────────────────────────────
fmt_currency() {
    printf "$%.2f" "$1"
}

fmt_delta() {
    local curr="$1" prev="$2"
    if (( $(echo "$prev == 0" | bc -l) )); then
        echo "N/A"
        return
    fi
    local pct
    pct=$(echo "scale=1; (($curr - $prev) / $prev) * 100" | bc -l)
    local sign=""
    if (( $(echo "$pct > 0" | bc -l) )); then
        sign="+"
    fi
    echo "${sign}${pct}%"
}

print_header() {
    local account_id
    account_id=$(aws sts get-caller-identity --query Account --output text 2>/dev/null || echo "unknown")

    echo "AWS Cost Reporter"
    echo "Account: $account_id"
    echo "Region:  $AWS_REGION"
    echo "Mode:    $RUN_MODE"
    echo "Time:    $(date -u +%Y-%m-%dT%H:%M:%SZ)"

    if [[ "$RUN_MODE" == "custom" ]]; then
        echo "Period:  $PERIOD_START → $PERIOD_END"
    fi
    echo ""
}

# ── Text table output ────────────────────────────────────────────────
output_text_table() {
    local -n curr_data=$1
    local -n prev_data=$2
    local label="SERVICE"
    case "$GROUP_BY" in
        LINKED_ACCOUNT) label="ACCOUNT" ;;
        TAG)            label="TAG" ;;
    esac
    local divider="──────────────────────────────────────────────────────────────────────"
    printf "  %-38s %-12s %-12s %s\n" "$label" "COST" "PREV" "DELTA"
    printf "  %s\n" "$divider"
    local total_curr=0 total_prev=0
    for key in "${!curr_data[@]}"; do
        local cost="${curr_data[$key]}" prev_cost="${prev_data[$key]:-0}"
        printf "  %-38s %-12s %-12s %s\n" \
            "$key" "$(fmt_currency "$cost")" "$(fmt_currency "$prev_cost")" "$(fmt_delta "$cost" "$prev_cost")"
        total_curr=$(echo "$total_curr + $cost" | bc -l)
        total_prev=$(echo "$total_prev + $prev_cost" | bc -l)
    done
    printf "  %s\n" "$divider"
    printf "  %-38s %-12s %-12s %s\n" \
        "TOTAL" "$(fmt_currency "$total_curr")" "$(fmt_currency "$total_prev")" "$(fmt_delta "$total_curr" "$total_prev")"
}

# ── CSV output ────────────────────────────────────────────────────────
output_csv() {
    local -n curr_data=$1
    local -n prev_data=$2
    local label="service"
    case "$GROUP_BY" in
        LINKED_ACCOUNT) label="account" ;;
        TAG)            label="tag" ;;
    esac
    echo "${label},cost,previous_cost,delta_pct"
    for key in "${!curr_data[@]}"; do
        local cost="${curr_data[$key]}" prev_cost="${prev_data[$key]:-0}" pct="0"
        if (( $(echo "$prev_cost != 0" | bc -l) )); then
            pct=$(echo "scale=2; (($cost - $prev_cost) / $prev_cost) * 100" | bc -l)
        fi
        echo "\"$key\",$cost,$prev_cost,$pct"
    done
}

# ── JSON output ───────────────────────────────────────────────────────
output_json() {
    local -n curr_data=$1
    local -n prev_data=$2
    local label="service"
    case "$GROUP_BY" in
        LINKED_ACCOUNT) label="account" ;;
        TAG)            label="tag" ;;
    esac
    local items=()
    for key in "${!curr_data[@]}"; do
        items+=("{\"${label}\": \"${key}\", \"cost\": ${curr_data[$key]}, \"previous_cost\": ${prev_data[$key]:-0}}")
    done
    local joined
    joined=$(printf ",%s" "${items[@]}")
    joined="${joined:1}"
    printf '{"mode":"%s","period_start":"%s","period_end":"%s","previous_start":"%s","previous_end":"%s","group_by":"%s","items":[%s]}\n' \
        "$RUN_MODE" "$PERIOD_START" "$PERIOD_END" "$PREV_START" "$PREV_END" "$GROUP_BY" "$joined"
}

# ── Render report ─────────────────────────────────────────────────────
render_report() {
    local curr_raw="$1" prev_raw="$2"

    # Parse into associative arrays
    declare -A curr_costs
    declare -A prev_costs

    while IFS=$'\t' read -r key amount; do
        [[ -z "$key" ]] && continue
        curr_costs["$key"]="$amount"
    done <<< "$(parse_costs "$curr_raw")"

    while IFS=$'\t' read -r key amount; do
        [[ -z "$key" ]] && continue
        prev_costs["$key"]="$amount"
    done <<< "$(parse_costs "$prev_raw")"

    # Ensure previous-only keys appear in current with 0
    for key in "${!prev_costs[@]}"; do
        if [[ -z "${curr_costs[$key]+x}" ]]; then
            curr_costs["$key"]="0"
        fi
    done

    case "$OUTPUT_FORMAT" in
        text)
            print_header
            local title="Cost Breakdown — ${PERIOD_START} → ${PERIOD_END}"
            echo "$title"
            output_text_table curr_costs prev_costs
            echo ""
            ;;
        csv)
            output_csv curr_costs prev_costs
            ;;
        json)
            output_json curr_costs prev_costs
            ;;
        *)
            die "Unknown format: $OUTPUT_FORMAT"
            ;;
    esac
}

# ── Email via SES ─────────────────────────────────────────────────────
send_email() {
    local report="$1" recipient="$2"

    if [[ -z "$SES_FROM_ADDRESS" ]]; then
        die "--email requires SES_FROM_ADDRESS to be set"
    fi

    local subject
    subject="AWS Cost Report — ${RUN_MODE} — $(today_utc)"

    log_info "Sending report to $recipient via SES..."

    local message
    message=$(jq -n \
        --arg from "$SES_FROM_ADDRESS" \
        --arg to "$recipient" \
        --arg subject "$subject" \
        --arg body "$report" \
        '{
            Source: $from,
            Destination: { ToAddresses: [$to] },
            Message: {
                Subject: { Data: $subject, Charset: "UTF-8" },
                Body: { Text: { Data: $body, Charset: "UTF-8" } }
            }
        }')

    aws ses send-email \
        --region "$AWS_REGION" \
        --cli-input-json "$message" \
        --output text >/dev/null

    log_info "Email sent to $recipient"
}

# ── Slack webhook ─────────────────────────────────────────────────────
send_slack() {
    local report="$1" webhook="$2"

    log_info "Posting report to Slack..."

    # Truncate for Slack message limits
    local max_len=3000
    local body="$report"
    if (( ${#body} > max_len )); then
        body="${body:0:$max_len}

... (truncated — full report exceeds Slack message limit)"
    fi

    local payload
    payload=$(jq -n --arg text "\`\`\`${body}\`\`\`" '{ text: $text }')

    local http_code
    http_code=$(curl -s -o /dev/null -w "%{http_code}" \
        -X POST \
        -H "Content-Type: application/json" \
        -d "$payload" \
        "$webhook")

    if [[ "$http_code" != "200" ]]; then
        log_error "Slack webhook returned HTTP $http_code"
        return 1
    fi

    log_info "Slack message posted"
}

# ── Usage ─────────────────────────────────────────────────────────────
usage() {
    cat <<EOF
Usage: $SCRIPT_NAME [MODE] [OPTIONS]

Modes:
  --daily                Today + yesterday cost breakdown
  --weekly               Last 7 days vs previous 7 days
  --monthly              Current month to date vs previous month
  --custom START END     Custom date range (YYYY-MM-DD)

Options:
  --group-by TYPE        SERVICE (default), TAG, LINKED_ACCOUNT
  --tag KEY=VALUE        Filter by cost allocation tag
  --format FORMAT        text (default), csv, json
  --email ADDRESS        Send report via SES (requires SES_FROM_ADDRESS)
  --slack WEBHOOK_URL    Post report to Slack webhook
  --verbose              Debug output
  --no-color             Disable colored output
  --help                 Show this help
EOF
    exit 0
}

# ── Argument parsing ─────────────────────────────────────────────────
parse_args() {
    while (( $# > 0 )); do
        case "$1" in
            --daily|--weekly|--monthly)
                RUN_MODE="${1#--}"; shift ;;
            --custom)
                RUN_MODE="custom"
                [[ $# -lt 3 ]] && die "--custom requires START and END dates"
                CUSTOM_START="$2"; CUSTOM_END="$3"
                validate_date "$CUSTOM_START"; validate_date "$CUSTOM_END"
                shift 3 ;;
            --group-by)
                [[ $# -lt 2 ]] && die "--group-by requires a value"
                GROUP_BY="$2"; shift 2 ;;
            --tag)
                [[ $# -lt 2 ]] && die "--tag requires KEY=VALUE"
                [[ "$2" != *"="* ]] && die "--tag value must be KEY=VALUE"
                COST_TAG_KEY="${2%%=*}"; COST_TAG_VALUE="${2#*=}"; shift 2 ;;
            --format)
                [[ $# -lt 2 ]] && die "--format requires a value"
                OUTPUT_FORMAT="$2"; shift 2 ;;
            --email)
                [[ $# -lt 2 ]] && die "--email requires an address"
                EMAIL_TO="$2"; shift 2 ;;
            --slack)
                [[ $# -lt 2 ]] && die "--slack requires a webhook URL"
                SLACK_URL="$2"; shift 2 ;;
            --verbose)  VERBOSE="true"; shift ;;
            --no-color) COLOR="never"; shift ;;
            --help|-h)  usage ;;
            *)          die "Unknown option: $1 (see --help)" ;;
        esac
    done

    if [[ -z "$RUN_MODE" ]]; then log_error "No mode specified"; echo ""; usage; exit 1; fi
    [[ -z "$SLACK_URL" && -n "$SLACK_WEBHOOK_URL" ]] && SLACK_URL="$SLACK_WEBHOOK_URL"

    case "$GROUP_BY" in
        SERVICE|TAG|LINKED_ACCOUNT) ;;
        *) die "Invalid --group-by: $GROUP_BY" ;;
    esac
    case "$OUTPUT_FORMAT" in
        text|csv|json) ;;
        *) die "Invalid --format: $OUTPUT_FORMAT" ;;
    esac
}

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

    START_TIME=$(date +%s)

    # Validate AWS credentials
    log_debug "Validating AWS credentials..."
    aws sts get-caller-identity --output text >/dev/null 2>&1 \
        || die "AWS credentials not configured or expired"

    compute_ranges

    log_info "Querying Cost Explorer ($RUN_MODE, group by $GROUP_BY)..."

    local curr_raw prev_raw
    curr_raw="$(query_costs "$PERIOD_START" "$PERIOD_END")"
    prev_raw="$(query_costs "$PREV_START" "$PREV_END")"

    if [[ -z "$curr_raw" ]]; then
        die "No cost data returned for $PERIOD_START → $PERIOD_END"
    fi

    local report
    report="$(render_report "$curr_raw" "$prev_raw")"

    # Output to stdout unless sending elsewhere exclusively
    echo "$report"

    # Email delivery
    if [[ -n "$EMAIL_TO" ]]; then
        send_email "$report" "$EMAIL_TO"
    fi

    # Slack delivery
    if [[ -n "$SLACK_URL" ]]; then
        send_slack "$report" "$SLACK_URL"
    fi

    local elapsed=$(( $(date +%s) - START_TIME ))
    log_info "Completed in ${elapsed}s"
}

main "$@"