#!/usr/bin/env bash

#########################################################################################
#### ec2-inventory-reporter.sh — AWS EC2 instance inventory and compliance report    ####
#### Instance metadata, uptime, cost estimates, tag compliance, SG audit             ####
#### Requires: bash 4+, aws-cli v2, jq                                               ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.00                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./ec2-inventory-reporter.sh                                                   ####
####   ./ec2-inventory-reporter.sh --all-regions --format csv                        ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
AWS_REGION="${AWS_REGION:-us-east-1}"
OUTPUT_FORMAT="${OUTPUT_FORMAT:-text}"
REQUIRED_TAGS="${REQUIRED_TAGS:-Name,Environment,Owner}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
SCAN_REGION="$AWS_REGION"
ALL_REGIONS="false"
FILTER_STATE=""
FILTER_TAG_KEY=""
FILTER_TAG_VALUE=""
FILTER_TYPE=""
TAG_CHECK="false"
SG_AUDIT="false"
START_TIME=""

# ── Colors ────────────────────────────────────────────────────────────
RED="" GREEN="" YELLOW="" BLUE="" BOLD="" DIM="" RESET=""
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        return
    fi
    if [[ "$COLOR" == "auto" && ! -t 1 ]]; then
        return
    fi
    RED="\033[0;31m"
    GREEN="\033[0;32m"
    YELLOW="\033[0;33m"
    BLUE="\033[0;34m"
    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")
    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
}

# ── Pricing table (us-east-1 on-demand Linux, $/hr) ──────────────────
declare -A PRICING=(
    # General purpose
    ["t3.nano"]=0.0052     ["t3.micro"]=0.0104    ["t3.small"]=0.0208
    ["t3.medium"]=0.0416   ["t3.large"]=0.0832    ["t3.xlarge"]=0.1664
    ["t3.2xlarge"]=0.3328
    ["t3a.nano"]=0.0047    ["t3a.micro"]=0.0094   ["t3a.small"]=0.0188
    ["t3a.medium"]=0.0376  ["t3a.large"]=0.0752   ["t3a.xlarge"]=0.1504
    ["t3a.2xlarge"]=0.3008
    ["m5.large"]=0.096     ["m5.xlarge"]=0.192    ["m5.2xlarge"]=0.384
    ["m5.4xlarge"]=0.768   ["m5.8xlarge"]=1.536
    ["m6i.large"]=0.096    ["m6i.xlarge"]=0.192   ["m6i.2xlarge"]=0.384
    ["m6i.4xlarge"]=0.768
    ["m7i.large"]=0.1008   ["m7i.xlarge"]=0.2016  ["m7i.2xlarge"]=0.4032
    # Compute optimized
    ["c5.large"]=0.085     ["c5.xlarge"]=0.17     ["c5.2xlarge"]=0.34
    ["c5.4xlarge"]=0.68    ["c5.9xlarge"]=1.53
    ["c6i.large"]=0.085    ["c6i.xlarge"]=0.17    ["c6i.2xlarge"]=0.34
    # Memory optimized
    ["r5.large"]=0.126     ["r5.xlarge"]=0.252    ["r5.2xlarge"]=0.504
    ["r5.4xlarge"]=1.008
    ["r6i.large"]=0.126    ["r6i.xlarge"]=0.252   ["r6i.2xlarge"]=0.504
    # Storage optimized
    ["i3.large"]=0.156     ["i3.xlarge"]=0.312    ["i3.2xlarge"]=0.624
    # Accelerated
    ["g4dn.xlarge"]=0.526  ["g4dn.2xlarge"]=0.752
    # Burstable previous gen
    ["t2.nano"]=0.0058     ["t2.micro"]=0.0116    ["t2.small"]=0.023
    ["t2.medium"]=0.0464   ["t2.large"]=0.0928
)

# ── Cost estimation ──────────────────────────────────────────────────
estimate_cost() {
    local instance_type="$1" state="$2"
    if [[ "$state" != "running" ]]; then
        echo "0.00"
        return
    fi
    local hourly="${PRICING[$instance_type]:-}"
    if [[ -z "$hourly" ]]; then
        echo "N/A"
        return
    fi
    printf "%.2f" "$(echo "$hourly * 730" | bc -l)"
}

# ── Uptime calculation ───────────────────────────────────────────────
format_uptime() {
    local launch_time="$1" state="$2"
    if [[ "$state" != "running" || -z "$launch_time" || "$launch_time" == "null" ]]; then
        echo "—"
        return
    fi

    local launch_epoch now_epoch diff_sec
    if date --version >/dev/null 2>&1; then
        launch_epoch=$(date -d "$launch_time" +%s 2>/dev/null) || { echo "—"; return; }
    else
        launch_epoch=$(date -j -f "%Y-%m-%dT%H:%M:%S" "${launch_time%%.*}" +%s 2>/dev/null) || { echo "—"; return; }
    fi
    now_epoch=$(date -u +%s)
    diff_sec=$(( now_epoch - launch_epoch ))

    if (( diff_sec < 0 )); then
        echo "—"
        return
    fi

    local days=$(( diff_sec / 86400 ))
    local hours=$(( (diff_sec % 86400) / 3600 ))
    local mins=$(( (diff_sec % 3600) / 60 ))
    printf "%dd %dh %dm" "$days" "$hours" "$mins"
}

# ── Tag compliance check ─────────────────────────────────────────────
check_tag_compliance() {
    local tags_json="$1"
    local missing=()

    IFS=',' read -ra required <<< "$REQUIRED_TAGS"
    for tag in "${required[@]}"; do
        tag=$(echo "$tag" | xargs)
        local found
        found=$(echo "$tags_json" | jq -r --arg key "$tag" '.[] | select(.Key == $key) | .Key' 2>/dev/null)
        if [[ -z "$found" ]]; then
            missing+=("$tag")
        fi
    done

    if (( ${#missing[@]} == 0 )); then
        echo "PASS"
    else
        echo "MISSING: ${missing[*]}"
    fi
}

# ── Security group audit ─────────────────────────────────────────────
audit_security_groups() {
    local region="$1"
    shift
    local sg_ids=("$@")
    local findings=()

    if (( ${#sg_ids[@]} == 0 )); then
        echo "—"
        return
    fi

    local sg_data
    sg_data=$(aws ec2 describe-security-groups \
        --region "$region" \
        --group-ids "${sg_ids[@]}" \
        --output json 2>/dev/null) || { echo "ERROR"; return; }

    local open_rules
    open_rules=$(echo "$sg_data" | jq -r '
        .SecurityGroups[].IpPermissions[] |
        select(
            (.IpRanges[]?.CidrIp == "0.0.0.0/0") or
            (.Ipv6Ranges[]?.CidrIpv6 == "::/0")
        ) |
        select(
            (.FromPort != 80 or .ToPort != 80) and
            (.FromPort != 443 or .ToPort != 443)
        ) |
        if .FromPort == .ToPort then
            "port \(.FromPort // "all")"
        elif .FromPort == -1 then
            "all ports"
        else
            "ports \(.FromPort)-\(.ToPort)"
        end
    ' 2>/dev/null)

    if [[ -z "$open_rules" ]]; then
        echo "OK"
    else
        local unique
        unique=$(echo "$open_rules" | sort -u | paste -sd ", " -)
        echo "OPEN: $unique"
    fi
}

# ── Query EC2 instances ──────────────────────────────────────────────
get_instances() {
    local region="$1"
    local filters=()

    if [[ -n "$FILTER_STATE" ]]; then
        filters+=("Name=instance-state-name,Values=$FILTER_STATE")
    fi
    if [[ -n "$FILTER_TAG_KEY" && -n "$FILTER_TAG_VALUE" ]]; then
        filters+=("Name=tag:$FILTER_TAG_KEY,Values=$FILTER_TAG_VALUE")
    fi
    if [[ -n "$FILTER_TYPE" ]]; then
        filters+=("Name=instance-type,Values=$FILTER_TYPE")
    fi

    local cmd=(
        aws ec2 describe-instances
        --region "$region"
        --output json
    )

    if (( ${#filters[@]} > 0 )); then
        cmd+=(--filters "${filters[@]}")
    fi

    log_debug "Running: ${cmd[*]}"

    local result=""
    local next_token=""

    while true; do
        local page_cmd=("${cmd[@]}")
        if [[ -n "$next_token" ]]; then
            page_cmd+=(--starting-token "$next_token")
        fi

        local page
        page=$("${page_cmd[@]}" 2>/dev/null) || { log_warn "Failed to query EC2 in $region"; echo "[]"; return; }

        local page_instances
        page_instances=$(echo "$page" | jq '[.Reservations[].Instances[]]')

        if [[ -z "$result" ]]; then
            result="$page_instances"
        else
            result=$(echo "$result $page_instances" | jq -s 'add')
        fi

        next_token=$(echo "$page" | jq -r '.NextToken // empty')
        if [[ -z "$next_token" ]]; then
            break
        fi
    done

    echo "$result"
}

# ── Get all enabled regions ──────────────────────────────────────────
get_all_regions() {
    aws ec2 describe-regions \
        --region "$AWS_REGION" \
        --query 'Regions[].RegionName' \
        --output text 2>/dev/null | tr '\t' '\n' | sort
}

# ── Text table output ────────────────────────────────────────────────
output_text() {
    local region="$1" instances_json="$2"
    local count
    count=$(echo "$instances_json" | jq 'length')

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

    echo "EC2 Inventory Reporter"
    echo "Account: $account_id"
    echo "Region:  $region"
    echo "Time:    $(date -u +%Y-%m-%dT%H:%M:%SZ)"
    echo "Instances: $count"
    echo ""

    if (( count == 0 )); then
        echo "  No instances found."
        echo ""
        return
    fi

    local divider="─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────"

    printf "  %-21s %-14s %-10s %-15s %-16s %-17s %-16s %s\n" \
        "INSTANCE ID" "TYPE" "STATE" "AZ" "PRIVATE IP" "PUBLIC IP" "UPTIME" "EST \$/MO"
    printf "  %s\n" "$divider"

    local total_cost=0 running=0 stopped=0 other=0
    local compliance_issues=0 sg_issues=0

    while IFS=$'\t' read -r iid itype istate iaz pip eip launch_time tags_json sg_json; do
        local uptime
        uptime=$(format_uptime "$launch_time" "$istate")

        local cost
        cost=$(estimate_cost "$itype" "$istate")

        [[ -z "$eip" || "$eip" == "null" ]] && eip="—"

        printf "  %-21s %-14s %-10s %-15s %-16s %-17s %-16s %s\n" \
            "$iid" "$itype" "$istate" "$iaz" "$pip" "$eip" "$uptime" "\$$cost"

        if [[ "$cost" != "N/A" ]]; then
            total_cost=$(echo "$total_cost + $cost" | bc -l)
        fi

        case "$istate" in
            running)   (( running++ )) ;;
            stopped)   (( stopped++ )) ;;
            *)         (( other++ )) ;;
        esac

        if [[ "$TAG_CHECK" == "true" ]]; then
            local compliance
            compliance=$(check_tag_compliance "$tags_json")
            if [[ "$compliance" != "PASS" ]]; then
                (( compliance_issues++ ))
                printf "  ${YELLOW}  ↳ Tag compliance: %s${RESET}\n" "$compliance"
            fi
        fi

        if [[ "$SG_AUDIT" == "true" && "$istate" == "running" ]]; then
            local sg_ids_list
            sg_ids_list=$(echo "$sg_json" | jq -r '.[].GroupId' 2>/dev/null)
            if [[ -n "$sg_ids_list" ]]; then
                local sg_arr=()
                while IFS= read -r sg; do
                    [[ -n "$sg" ]] && sg_arr+=("$sg")
                done <<< "$sg_ids_list"
                local sg_result
                sg_result=$(audit_security_groups "$region" "${sg_arr[@]}")
                if [[ "$sg_result" != "OK" && "$sg_result" != "—" ]]; then
                    (( sg_issues++ ))
                    printf "  ${RED}  ↳ SG audit: %s${RESET}\n" "$sg_result"
                fi
            fi
        fi

    done < <(echo "$instances_json" | jq -r '
        .[] |
        [
            .InstanceId,
            .InstanceType,
            (.State.Name),
            (.Placement.AvailabilityZone),
            (.PrivateIpAddress // "—"),
            (.PublicIpAddress // "null"),
            (.LaunchTime // "null"),
            (.Tags // [] | tojson),
            (.SecurityGroups // [] | tojson)
        ] | @tsv
    ')

    printf "  %s\n" "$divider"

    local summary="TOTAL: $count instances"
    local parts=()
    (( running > 0 )) && parts+=("$running running")
    (( stopped > 0 )) && parts+=("$stopped stopped")
    (( other > 0 ))   && parts+=("$other other")
    if (( ${#parts[@]} > 0 )); then
        local joined
        joined=$(printf ", %s" "${parts[@]}")
        summary+=" (${joined:2})"
    fi
    printf "  %-70s Estimated monthly cost: \$%.2f\n" "$summary" "$total_cost"

    if [[ "$TAG_CHECK" == "true" ]]; then
        echo ""
        if (( compliance_issues > 0 )); then
            printf "  ${YELLOW}Tag compliance issues: %d instance(s) missing required tags${RESET}\n" "$compliance_issues"
        else
            printf "  ${GREEN}Tag compliance: all instances have required tags${RESET}\n"
        fi
    fi

    if [[ "$SG_AUDIT" == "true" ]]; then
        if (( sg_issues > 0 )); then
            printf "  ${RED}Security group issues: %d instance(s) with overly permissive rules${RESET}\n" "$sg_issues"
        else
            printf "  ${GREEN}Security groups: no overly permissive rules found${RESET}\n"
        fi
    fi

    echo ""
}

# ── CSV output ────────────────────────────────────────────────────────
output_csv() {
    local region="$1" instances_json="$2"
    local count
    count=$(echo "$instances_json" | jq 'length')

    local header="instance_id,type,state,az,private_ip,public_ip,launch_time,uptime,est_monthly_cost"
    if [[ "$TAG_CHECK" == "true" ]]; then
        header+=",tag_compliance"
    fi
    if [[ "$SG_AUDIT" == "true" ]]; then
        header+=",sg_audit"
    fi
    header+=",region"
    echo "$header"

    if (( count == 0 )); then
        return
    fi

    while IFS=$'\t' read -r iid itype istate iaz pip eip launch_time tags_json sg_json; do
        local uptime
        uptime=$(format_uptime "$launch_time" "$istate")

        local cost
        cost=$(estimate_cost "$itype" "$istate")

        [[ -z "$eip" || "$eip" == "null" ]] && eip=""

        local line="$iid,$itype,$istate,$iaz,$pip,$eip,$launch_time,\"$uptime\",$cost"

        if [[ "$TAG_CHECK" == "true" ]]; then
            local compliance
            compliance=$(check_tag_compliance "$tags_json")
            line+=",\"$compliance\""
        fi

        if [[ "$SG_AUDIT" == "true" ]]; then
            local sg_ids_list
            sg_ids_list=$(echo "$sg_json" | jq -r '.[].GroupId' 2>/dev/null)
            if [[ -n "$sg_ids_list" && "$istate" == "running" ]]; then
                local sg_arr=()
                while IFS= read -r sg; do
                    [[ -n "$sg" ]] && sg_arr+=("$sg")
                done <<< "$sg_ids_list"
                local sg_result
                sg_result=$(audit_security_groups "$region" "${sg_arr[@]}")
                line+=",\"$sg_result\""
            else
                line+=",\"—\""
            fi
        fi

        line+=",$region"
        echo "$line"

    done < <(echo "$instances_json" | jq -r '
        .[] |
        [
            .InstanceId,
            .InstanceType,
            (.State.Name),
            (.Placement.AvailabilityZone),
            (.PrivateIpAddress // "—"),
            (.PublicIpAddress // "null"),
            (.LaunchTime // "null"),
            (.Tags // [] | tojson),
            (.SecurityGroups // [] | tojson)
        ] | @tsv
    ')
}

# ── JSON output ───────────────────────────────────────────────────────
output_json() {
    local region="$1" instances_json="$2"
    local count
    count=$(echo "$instances_json" | jq 'length')

    local items="[]"

    if (( count > 0 )); then
        items=$(echo "$instances_json" | jq --arg region "$region" '[
            .[] | {
                instance_id: .InstanceId,
                type: .InstanceType,
                state: .State.Name,
                az: .Placement.AvailabilityZone,
                private_ip: (.PrivateIpAddress // null),
                public_ip: (.PublicIpAddress // null),
                launch_time: (.LaunchTime // null),
                ami_id: (.ImageId // null),
                vpc_id: (.VpcId // null),
                key_name: (.KeyName // null),
                tags: (.Tags // []),
                security_groups: (.SecurityGroups // []),
                region: $region
            }
        ]')
    fi

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

    jq -n \
        --arg account "$account_id" \
        --arg region "$region" \
        --arg time "$(date -u +%Y-%m-%dT%H:%M:%SZ)" \
        --argjson count "$count" \
        --argjson instances "$items" \
        '{
            account: $account,
            region: $region,
            time: $time,
            instance_count: $count,
            instances: $instances
        }'
}

# ── Process a single region ──────────────────────────────────────────
process_region() {
    local region="$1"
    local csv_header_printed="$2"

    log_debug "Querying region: $region"

    local instances_json
    instances_json=$(get_instances "$region")

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

    if (( count == 0 )) && [[ "$ALL_REGIONS" == "true" ]]; then
        log_debug "No instances in $region, skipping"
        return
    fi

    case "$OUTPUT_FORMAT" in
        text)
            output_text "$region" "$instances_json"
            ;;
        csv)
            if [[ "$csv_header_printed" == "false" ]]; then
                output_csv "$region" "$instances_json"
            else
                output_csv "$region" "$instances_json" | tail -n +2
            fi
            ;;
        json)
            output_json "$region" "$instances_json"
            ;;
    esac
}

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

Options:
  --region REGION       AWS region (default: \$AWS_REGION or us-east-1)
  --all-regions         Scan all enabled regions
  --state STATE         Filter by state (running, stopped, terminated, etc.)
  --tag KEY=VALUE       Filter by tag
  --type TYPE           Filter by instance type (e.g., t3.micro, m5.*)
  --format FORMAT       Output: text (default), csv, json
  --tag-check           Check tag compliance (requires: $REQUIRED_TAGS)
  --sg-audit            Audit security groups for overly permissive rules
  --verbose             Debug output
  --no-color            Disable colored output
  --help                Show this help

Environment Variables:
  AWS_PROFILE           AWS CLI profile
  AWS_REGION            Default region (default: us-east-1)
  REQUIRED_TAGS         Comma-separated required tags (default: Name,Environment,Owner)
  OUTPUT_FORMAT         Default output format (default: text)

Examples:
  $SCRIPT_NAME                                    # All instances, default region
  $SCRIPT_NAME --state running                    # Running instances only
  $SCRIPT_NAME --all-regions --format csv         # All regions, CSV output
  $SCRIPT_NAME --tag-check --sg-audit             # Full compliance audit
EOF
    exit 0
}

# ── Argument parsing ─────────────────────────────────────────────────
parse_args() {
    while (( $# > 0 )); do
        case "$1" in
            --region)
                [[ $# -lt 2 ]] && die "--region requires a value"
                SCAN_REGION="$2"; shift 2 ;;
            --all-regions)
                ALL_REGIONS="true"; shift ;;
            --state)
                [[ $# -lt 2 ]] && die "--state requires a value"
                FILTER_STATE="$2"; shift 2 ;;
            --tag)
                [[ $# -lt 2 ]] && die "--tag requires KEY=VALUE"
                [[ "$2" != *"="* ]] && die "--tag value must be KEY=VALUE"
                FILTER_TAG_KEY="${2%%=*}"; FILTER_TAG_VALUE="${2#*=}"; shift 2 ;;
            --type)
                [[ $# -lt 2 ]] && die "--type requires a value"
                FILTER_TYPE="$2"; shift 2 ;;
            --format)
                [[ $# -lt 2 ]] && die "--format requires a value"
                OUTPUT_FORMAT="$2"; shift 2 ;;
            --tag-check)
                TAG_CHECK="true"; shift ;;
            --sg-audit)
                SG_AUDIT="true"; shift ;;
            --verbose)
                VERBOSE="true"; shift ;;
            --no-color)
                COLOR="never"; shift ;;
            --help|-h)
                usage ;;
            *)
                die "Unknown option: $1 (see --help)" ;;
        esac
    done

    case "$OUTPUT_FORMAT" in
        text|csv|json) ;;
        *) die "Invalid --format: $OUTPUT_FORMAT (expected text, csv, json)" ;;
    esac
}

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

    START_TIME=$(date +%s)

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

    if [[ "$ALL_REGIONS" == "true" ]]; then
        log_info "Scanning all enabled regions..."
        local regions
        regions=$(get_all_regions)

        if [[ -z "$regions" ]]; then
            die "Failed to retrieve region list"
        fi

        local csv_header="false"
        local json_first="true"

        if [[ "$OUTPUT_FORMAT" == "json" ]]; then
            echo "["
        fi

        while IFS= read -r region; do
            [[ -z "$region" ]] && continue
            log_info "Scanning $region..."

            if [[ "$OUTPUT_FORMAT" == "json" ]]; then
                if [[ "$json_first" == "true" ]]; then
                    json_first="false"
                else
                    echo ","
                fi
            fi

            process_region "$region" "$csv_header"
            csv_header="true"
        done <<< "$regions"

        if [[ "$OUTPUT_FORMAT" == "json" ]]; then
            echo "]"
        fi
    else
        log_info "Scanning region: $SCAN_REGION"
        process_region "$SCAN_REGION" "false"
    fi

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

main "$@"