#!/usr/bin/env bash

#########################################################################################
#### service-dependency-mapper.sh — Map systemd service dependencies                 ####
#### Shows what depends on a service and what it requires                             ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.00                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./service-dependency-mapper.sh nginx                                          ####
####   ./service-dependency-mapper.sh --all --format tree                            ####
####   ./service-dependency-mapper.sh postgresql --reverse                           ####
####   ./service-dependency-mapper.sh --critical                                     ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
SERVICE=""
MODE="forward"
MAP_ALL="false"
CRITICAL="false"
DEPTH=3
FORMAT="tree"
COLOR="auto"

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME

# ── Colors ────────────────────────────────────────────────────────────
RED="" GREEN="" YELLOW="" CYAN="" BOLD="" DIM="" RESET=""

setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BOLD="" DIM="" CYAN="" RESET=""
        return
    fi
    if [[ "$COLOR" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        CYAN='\033[0;36m'
        BOLD='\033[1m'
        DIM='\033[2m'
        RESET='\033[0m'
    else
        RED="" GREEN="" YELLOW="" BOLD="" DIM="" CYAN="" RESET=""
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
err()  { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
warn() { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }

# ══════════════════════════════════════════════════════════════════════
#                       USAGE
# ══════════════════════════════════════════════════════════════════════

usage() {
    cat <<EOF
${SCRIPT_NAME} — Map systemd service dependencies

USAGE:
  ${SCRIPT_NAME} SERVICE [OPTIONS]
  ${SCRIPT_NAME} --all [OPTIONS]
  ${SCRIPT_NAME} --critical [OPTIONS]

OPTIONS:
  --forward             Show what this service requires (default)
  --reverse             Show what depends ON this service
  --all                 Map all active services
  --critical            Show services with the most reverse dependencies
  --depth N             Depth limit for dependency traversal (default: 3)
  --format FORMAT       Output format: tree, flat, dot, json (default: tree)
  --no-color            Disable colored output
  --help                Show this help

EXAMPLES:
  # Show what nginx requires
  ./${SCRIPT_NAME} nginx

  # Show what breaks if postgresql stops
  ./${SCRIPT_NAME} postgresql --reverse

  # Show all active service dependencies as a tree
  ./${SCRIPT_NAME} --all --format tree

  # Find highest-impact services
  ./${SCRIPT_NAME} --critical

  # Generate Graphviz DOT output
  ./${SCRIPT_NAME} nginx --format dot | dot -Tpng -o nginx-deps.png

  # Shallow scan
  ./${SCRIPT_NAME} nginx --depth 1
EOF
}

# ══════════════════════════════════════════════════════════════════════
#                       ARGUMENT PARSING
# ══════════════════════════════════════════════════════════════════════

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --forward)
                MODE="forward"; shift ;;
            --reverse)
                MODE="reverse"; shift ;;
            --all)
                MAP_ALL="true"; shift ;;
            --critical)
                CRITICAL="true"; shift ;;
            --depth)
                DEPTH="$2"; shift 2 ;;
            --format)
                FORMAT="$2"; shift 2 ;;
            --no-color)
                COLOR="never"; shift ;;
            --help|-h)
                usage; exit 0 ;;
            -*)
                err "Unknown option: $1"; usage; exit 1 ;;
            *)
                SERVICE="$1"; shift ;;
        esac
    done

    if [[ "$MAP_ALL" == "false" && "$CRITICAL" == "false" && -z "$SERVICE" ]]; then
        err "No service specified. Use --all, --critical, or provide a service name."
        usage
        exit 1
    fi

    # Validate format
    case "$FORMAT" in
        tree|flat|dot|json) ;;
        *) err "Unknown format: $FORMAT"; exit 1 ;;
    esac

    # Validate depth is a positive integer
    if ! [[ "$DEPTH" =~ ^[0-9]+$ ]] || [[ "$DEPTH" -lt 1 ]]; then
        err "Depth must be a positive integer"
        exit 1
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                       HELPERS
# ══════════════════════════════════════════════════════════════════════

# Normalize unit name — append .service if no suffix present
normalize_unit() {
    local unit="$1"
    if [[ "$unit" != *.* ]]; then
        echo "${unit}.service"
    else
        echo "$unit"
    fi
}

# Get active state of a unit
get_active_state() {
    systemctl show -p ActiveState --value "$1" 2>/dev/null || echo "unknown"
}

# Get load state of a unit
get_load_state() {
    systemctl show -p LoadState --value "$1" 2>/dev/null || echo "unknown"
}

# Get unit description
get_description() {
    systemctl show -p Description --value "$1" 2>/dev/null || echo ""
}

# Color a unit name based on its active state
color_unit() {
    local unit="$1"
    local state
    state="$(get_active_state "$unit")"
    case "$state" in
        active)       echo -e "${GREEN}${unit}${RESET}" ;;
        failed)       echo -e "${RED}${unit}${RESET}" ;;
        inactive)     echo -e "${YELLOW}${unit}${RESET}" ;;
        *)            echo -e "${DIM}${unit}${RESET}" ;;
    esac
}

# Get list of all active units (services, sockets, timers, targets)
get_active_units() {
    systemctl list-units --type=service,socket,timer,target --state=active \
        --no-pager --no-legend --plain 2>/dev/null | awk '{print $1}'
}

# ══════════════════════════════════════════════════════════════════════
#                       DEPENDENCY COLLECTION
# ══════════════════════════════════════════════════════════════════════

# Collect dependencies using systemctl list-dependencies
# Args: unit, direction (forward|reverse), max_depth
# Outputs lines: "depth unit" (tab-indented raw from systemctl, parsed)
collect_deps() {
    local unit="$1"
    local direction="$2"
    local max_depth="$3"
    local reverse_flag=""

    [[ "$direction" == "reverse" ]] && reverse_flag="--reverse"

    # systemctl list-dependencies outputs an indented tree with unicode chars
    # We parse the depth from leading whitespace and extract the unit name
    systemctl list-dependencies "$unit" $reverse_flag --no-pager --plain 2>/dev/null | \
        tail -n +2 | while IFS= read -r line; do
            # Count leading spaces (each level = 2 spaces)
            local stripped="${line#"${line%%[! ]*}"}"
            local spaces=$(( ${#line} - ${#stripped} ))
            local depth=$(( spaces / 2 ))

            # Respect depth limit
            if [[ $depth -ge $max_depth ]]; then
                continue
            fi

            # Clean unit name (remove any tree-drawing characters)
            local dep_unit
            dep_unit="$(echo "$stripped" | sed 's/[●○]//g' | xargs)"

            if [[ -n "$dep_unit" ]]; then
                echo "${depth} ${dep_unit}"
            fi
        done
}

# ══════════════════════════════════════════════════════════════════════
#                       SERVICE INFO
# ══════════════════════════════════════════════════════════════════════

print_service_info() {
    local unit="$1"
    local active_state load_state description

    active_state="$(get_active_state "$unit")"
    load_state="$(get_load_state "$unit")"
    description="$(get_description "$unit")"

    echo ""
    echo -e "${BOLD}Service Info${RESET}"
    echo "  ──────────────────────────────────────────────────────────────"
    echo -e "  Name:          ${BOLD}${unit}${RESET}"
    echo -e "  Description:   ${description}"
    echo -e "  Load State:    ${load_state}"

    case "$active_state" in
        active)   echo -e "  Active State:  ${GREEN}${active_state}${RESET}" ;;
        failed)   echo -e "  Active State:  ${RED}${active_state}${RESET}" ;;
        inactive) echo -e "  Active State:  ${YELLOW}${active_state}${RESET}" ;;
        *)        echo -e "  Active State:  ${DIM}${active_state}${RESET}" ;;
    esac
    echo ""
}

# ══════════════════════════════════════════════════════════════════════
#                       OUTPUT FORMATTERS
# ══════════════════════════════════════════════════════════════════════

# ── Tree output ───────────────────────────────────────────────────────
output_tree() {
    local unit="$1"
    local direction="$2"
    local deps

    deps="$(collect_deps "$unit" "$direction" "$DEPTH")"

    if [[ "$direction" == "forward" ]]; then
        echo -e "${BOLD}Forward Dependencies${RESET} (what ${unit} requires):"
    else
        echo -e "${BOLD}Reverse Dependencies${RESET} (what depends on ${unit}):"
    fi
    echo "  ──────────────────────────────────────────────────────────────"

    if [[ -z "$deps" ]]; then
        echo "  (none)"
        echo ""
        return
    fi

    while IFS= read -r line; do
        local depth="${line%% *}"
        local dep_unit="${line#* }"
        local indent=""
        for (( i=0; i<depth; i++ )); do
            indent+="  "
        done
        local colored
        colored="$(color_unit "$dep_unit")"
        echo -e "  ${indent}├── ${colored}"
    done <<< "$deps"
    echo ""
}

# ── Flat output ───────────────────────────────────────────────────────
output_flat() {
    local unit="$1"
    local direction="$2"
    local deps

    deps="$(collect_deps "$unit" "$direction" "$DEPTH")"

    if [[ -z "$deps" ]]; then
        return
    fi

    while IFS= read -r line; do
        local dep_unit="${line#* }"
        echo "$dep_unit"
    done <<< "$deps" | sort -u
}

# ── DOT output ────────────────────────────────────────────────────────
output_dot() {
    local unit="$1"
    local direction="$2"
    local deps

    deps="$(collect_deps "$unit" "$direction" "$DEPTH")"

    echo "digraph dependencies {"
    echo "  rankdir=LR;"
    echo "  node [shape=box, fontname=\"monospace\", fontsize=10];"
    echo "  edge [color=\"#555555\"];"
    echo ""

    # Color the root node
    local root_state
    root_state="$(get_active_state "$unit")"
    local root_color="#4CAF50"
    [[ "$root_state" == "failed" ]] && root_color="#F44336"
    [[ "$root_state" == "inactive" ]] && root_color="#FF9800"
    echo "  \"${unit}\" [style=filled, fillcolor=\"${root_color}\", fontcolor=white];"
    echo ""

    if [[ -n "$deps" ]]; then
        while IFS= read -r line; do
            local dep_unit="${line#* }"
            if [[ "$direction" == "forward" ]]; then
                echo "  \"${unit}\" -> \"${dep_unit}\";"
            else
                echo "  \"${dep_unit}\" -> \"${unit}\";"
            fi
        done <<< "$deps"
    fi

    echo "}"
}

# ── JSON output ───────────────────────────────────────────────────────
output_json() {
    local unit="$1"
    local direction="$2"
    local deps active_state load_state description

    deps="$(collect_deps "$unit" "$direction" "$DEPTH")"
    active_state="$(get_active_state "$unit")"
    load_state="$(get_load_state "$unit")"
    description="$(get_description "$unit")"

    echo "{"
    echo "  \"unit\": \"${unit}\","
    echo "  \"active_state\": \"${active_state}\","
    echo "  \"load_state\": \"${load_state}\","
    echo "  \"description\": \"${description}\","
    echo "  \"direction\": \"${direction}\","
    echo "  \"depth_limit\": ${DEPTH},"
    echo "  \"dependencies\": ["

    local first="true"
    if [[ -n "$deps" ]]; then
        while IFS= read -r line; do
            local depth="${line%% *}"
            local dep_unit="${line#* }"
            local dep_state
            dep_state="$(get_active_state "$dep_unit")"

            if [[ "$first" == "true" ]]; then
                first="false"
            else
                echo ","
            fi
            printf '    {"unit": "%s", "depth": %s, "active_state": "%s"}' \
                "$dep_unit" "$depth" "$dep_state"
        done <<< "$deps"
        echo ""
    fi

    echo "  ]"
    echo "}"
}

# ══════════════════════════════════════════════════════════════════════
#                       IMPACT SUMMARY
# ══════════════════════════════════════════════════════════════════════

print_impact_summary() {
    local unit="$1"
    local reverse_deps

    reverse_deps="$(collect_deps "$unit" "reverse" "$DEPTH")"
    local count=0

    if [[ -n "$reverse_deps" ]]; then
        count="$(echo "$reverse_deps" | wc -l)"
    fi

    echo -e "${BOLD}Impact Summary${RESET}"
    echo "  ──────────────────────────────────────────────────────────────"
    if [[ $count -eq 0 ]]; then
        echo -e "  Stopping ${BOLD}${unit}${RESET} would affect ${GREEN}0${RESET} other units."
    elif [[ $count -le 3 ]]; then
        echo -e "  Stopping ${BOLD}${unit}${RESET} would affect ${YELLOW}${count}${RESET} other unit(s)."
    else
        echo -e "  Stopping ${BOLD}${unit}${RESET} would affect ${RED}${count}${RESET} other unit(s)."
    fi
    echo ""
}

# ══════════════════════════════════════════════════════════════════════
#                       CRITICAL MODE
# ══════════════════════════════════════════════════════════════════════

run_critical() {
    echo ""
    echo -e "${BOLD}Critical Services — ranked by reverse dependency count${RESET}"
    echo -e "${DIM}(services with the most units depending on them)${RESET}"
    echo "  ══════════════════════════════════════════════════════════════"
    echo ""

    local units
    units="$(get_active_units)"

    if [[ -z "$units" ]]; then
        err "No active units found"
        exit 1
    fi

    # Collect reverse dep counts
    local -a results=()
    while IFS= read -r unit; do
        local reverse_deps count=0
        reverse_deps="$(collect_deps "$unit" "reverse" "$DEPTH" 2>/dev/null)"
        if [[ -n "$reverse_deps" ]]; then
            count="$(echo "$reverse_deps" | wc -l)"
        fi
        if [[ $count -gt 0 ]]; then
            results+=("${count} ${unit}")
        fi
    done <<< "$units"

    # Sort by count descending, show top 20
    printf '%s\n' "${results[@]}" | sort -rn | head -20 | while IFS= read -r line; do
        local count="${line%% *}"
        local unit="${line#* }"
        local colored
        colored="$(color_unit "$unit")"

        if [[ $count -ge 10 ]]; then
            printf "  ${RED}%4d${RESET}  %b\n" "$count" "$colored"
        elif [[ $count -ge 5 ]]; then
            printf "  ${YELLOW}%4d${RESET}  %b\n" "$count" "$colored"
        else
            printf "  ${GREEN}%4d${RESET}  %b\n" "$count" "$colored"
        fi
    done

    echo ""
    echo -e "${DIM}  Top 20 shown. Higher count = more impact if stopped.${RESET}"
    echo ""
}

# ══════════════════════════════════════════════════════════════════════
#                       ALL MODE
# ══════════════════════════════════════════════════════════════════════

run_all() {
    local units
    units="$(get_active_units)"

    if [[ -z "$units" ]]; then
        err "No active units found"
        exit 1
    fi

    echo ""
    echo -e "${BOLD}All Active Service Dependencies${RESET} (${MODE} / depth ${DEPTH})"
    echo "  ══════════════════════════════════════════════════════════════"
    echo ""

    while IFS= read -r unit; do
        case "$FORMAT" in
            tree)
                echo -e "${CYAN}■ ${unit}${RESET}"
                output_tree "$unit" "$MODE"
                ;;
            flat)
                echo -e "# ${unit}"
                output_flat "$unit" "$MODE"
                echo ""
                ;;
            dot)
                output_dot "$unit" "$MODE"
                ;;
            json)
                output_json "$unit" "$MODE"
                ;;
        esac
    done <<< "$units"
}

# ══════════════════════════════════════════════════════════════════════
#                       SINGLE SERVICE MODE
# ══════════════════════════════════════════════════════════════════════

run_single() {
    local unit
    unit="$(normalize_unit "$SERVICE")"

    # Verify the unit exists
    local load_state
    load_state="$(get_load_state "$unit")"
    if [[ "$load_state" == "not-found" ]]; then
        err "Unit '${unit}' not found"
        exit 1
    fi

    case "$FORMAT" in
        tree)
            echo ""
            echo -e "${BOLD}Service Dependency Map — ${unit}${RESET}"
            echo -e "${DIM}$(date -u '+%Y-%m-%d %H:%M:%S UTC')${RESET}"

            print_service_info "$unit"

            if [[ "$MODE" == "forward" || "$MODE" == "both" ]]; then
                output_tree "$unit" "forward"
            fi
            if [[ "$MODE" == "reverse" || "$MODE" == "both" ]]; then
                output_tree "$unit" "reverse"
            fi

            # Always show both directions for single-service tree mode
            if [[ "$MODE" == "forward" ]]; then
                output_tree "$unit" "reverse"
            elif [[ "$MODE" == "reverse" ]]; then
                output_tree "$unit" "forward"
            fi

            print_impact_summary "$unit"
            ;;
        flat)
            output_flat "$unit" "$MODE"
            ;;
        dot)
            output_dot "$unit" "$MODE"
            ;;
        json)
            output_json "$unit" "$MODE"
            ;;
    esac
}

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

main() {
    parse_args "$@"
    setup_colors

    if [[ "$CRITICAL" == "true" ]]; then
        run_critical
    elif [[ "$MAP_ALL" == "true" ]]; then
        run_all
    else
        run_single
    fi
}

main "$@"