#!/usr/bin/env bash

#########################################################################################
#### azure-cost-reporter.sh — Azure cost breakdown by service, resource group, or    ####
#### tag. Supports email, Slack webhooks, CSV/JSON export, period comparison         ####
#### Requires: bash 4+, az CLI, jq, curl                                             ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.00                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./azure-cost-reporter.sh --daily                                              ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
SUBSCRIPTION="${SUBSCRIPTION:-}"
GROUP_BY="${GROUP_BY:-SERVICE}"
OUTPUT_FORMAT="${OUTPUT_FORMAT:-text}"
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=""
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 az   >/dev/null 2>&1 || missing+=("az-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"
}

# ── Resolve subscription ─────────────────────────────────────────────
resolve_subscription() {
    if [[ -n "$SUBSCRIPTION" ]]; then
        log_debug "Using subscription: $SUBSCRIPTION"
        return
    fi
    SUBSCRIPTION=$(az account show --query 'id' -o tsv 2>/dev/null) \
        || die "Cannot determine subscription. Use --subscription or az account set."
    log_debug "Resolved subscription: $SUBSCRIPTION"
}

# ── Build Cost Management query payload ──────────────────────────────
build_query_payload() {
    local start="$1" end="$2"
    local grouping_name grouping_type

    case "$GROUP_BY" in
        SERVICE)
            grouping_name="ServiceName"
            grouping_type="Dimension"
            ;;
        RESOURCE_GROUP)
            grouping_name="ResourceGroupName"
            grouping_type="Dimension"
            ;;
        TAG)
            if [[ -z "$COST_TAG_KEY" ]]; then
                die "--group-by TAG requires --tag KEY=VALUE"
            fi
            grouping_name="$COST_TAG_KEY"
            grouping_type="TagKey"
            ;;
        *)
            die "Invalid --group-by value: $GROUP_BY (expected SERVICE, RESOURCE_GROUP, or TAG)"
            ;;
    esac

    local filter_block="{}"
    if [[ -n "$COST_TAG_KEY" && -n "$COST_TAG_VALUE" && "$GROUP_BY" != "TAG" ]]; then
        filter_block=$(jq -n \
            --arg key "$COST_TAG_KEY" \
            --arg val "$COST_TAG_VALUE" \
            '{
                "Tags": {
                    "Name": $key,
                    "Operator": "In",
                    "Values": [$val]
                }
            }')
    fi

    local payload
    if [[ "$filter_block" == "{}" ]]; then
        payload=$(jq -n \
            --arg start "$start" \
            --arg end "$end" \
            --arg gname "$grouping_name" \
            --arg gtype "$grouping_type" \
            '{
                "type": "ActualCost",
                "dataSet": {
                    "granularity": "None",
                    "aggregation": {
                        "totalCost": {
                            "name": "Cost",
                            "function": "Sum"
                        }
                    },
                    "grouping": [
                        {
                            "type": $gtype,
                            "name": $gname
                        }
                    ]
                },
                "timeframe": "Custom",
                "timePeriod": {
                    "from": $start,
                    "to": $end
                }
            }')
    else
        payload=$(jq -n \
            --arg start "$start" \
            --arg end "$end" \
            --arg gname "$grouping_name" \
            --arg gtype "$grouping_type" \
            --argjson filter "$filter_block" \
            '{
                "type": "ActualCost",
                "dataSet": {
                    "granularity": "None",
                    "aggregation": {
                        "totalCost": {
                            "name": "Cost",
                            "function": "Sum"
                        }
                    },
                    "grouping": [
                        {
                            "type": $gtype,
                            "name": $gname
                        }
                    ],
                    "filter": $filter
                },
                "timeframe": "Custom",
                "timePeriod": {
                    "from": $start,
                    "to": $end
                }
            }')
    fi

    echo "$payload"
}

# ── Query Cost Management API ────────────────────────────────────────
query_costs() {
    local start="$1" end="$2"
    local payload
    payload="$(build_query_payload "$start" "$end")"

    local scope="/subscriptions/${SUBSCRIPTION}"
    local api_url="https://management.azure.com${scope}/providers/Microsoft.CostManagement/query?api-version=2023-11-01"

    log_debug "Querying: $api_url"
    log_debug "Payload: $payload"

    az rest \
        --method post \
        --url "$api_url" \
        --body "$payload" \
        --output json 2>/dev/null
}

# ── Parse cost data ──────────────────────────────────────────────────
parse_costs() {
    local raw="$1"
    echo "$raw" | jq -r '
        .properties.rows // [] |
        map({
            key: .[1],
            amount: (.[0] | 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 sub_name
    sub_name=$(az account show --query 'name' -o tsv 2>/dev/null || echo "unknown")

    echo "Azure Cost Reporter"
    echo "Subscription: $sub_name ($SUBSCRIPTION)"
    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
        RESOURCE_GROUP) label="RESOURCE_GROUP" ;;
        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
        RESOURCE_GROUP) label="resource_group" ;;
        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
        RESOURCE_GROUP) label="resource_group" ;;
        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
}

# ── 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                Yesterday vs day before 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), RESOURCE_GROUP, TAG
  --tag KEY=VALUE        Filter by cost allocation tag
  --subscription SUB     Azure subscription ID (default: current az account)
  --format FORMAT        text (default), csv, json
  --slack WEBHOOK_URL    Post report to Slack webhook
  --verbose              Debug output
  --no-color             Disable colored output
  --help                 Show this help

Environment:
  SUBSCRIPTION           Default subscription ID
  SLACK_WEBHOOK_URL      Default Slack webhook URL
  COLOR                  auto (default), always, never

Notes:
  Email delivery is not built in. Pipe output to sendmail or use an
  SMTP relay:  ./azure-cost-reporter.sh --daily | mail -s "Azure Costs" user@example.com
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 ;;
            --subscription)
                [[ $# -lt 2 ]] && die "--subscription requires a value"
                SUBSCRIPTION="$2"; shift 2 ;;
            --format)
                [[ $# -lt 2 ]] && die "--format requires a value"
                OUTPUT_FORMAT="$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|RESOURCE_GROUP|TAG) ;;
        *) 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 Azure credentials
    log_debug "Validating Azure credentials..."
    az account show >/dev/null 2>&1 \
        || die "Azure credentials not configured — run 'az login' first"

    resolve_subscription
    compute_ranges

    log_info "Querying Cost Management ($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
    echo "$report"

    # 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 "$@"