#!/usr/bin/env bash
# incusos-health -- Inspect IncusOS node health via the /os/1.0 REST API
#
# Queries one or more IncusOS cluster nodes for status, storage, security,
# services, network, and update information.
#
# Usage: incusos-health [OPTIONS] [ACTION]
# Run 'incusos-health --help' for full usage information.

set -euo pipefail

readonly VERSION="0.1.0"
readonly SCRIPT_NAME="incusos-health"

# ---------------------------------------------------------------------------
# Defaults
# ---------------------------------------------------------------------------

REMOTE="oc-node-01"
NODE=""
JSON_OUTPUT=false
DRY_RUN=false
VERBOSE=false
QUIET=false
ACTION=""

# Services to check (IncusOS built-in services)
readonly SERVICES=(iscsi lvm multipath nvme ovn tailscale usbip)

# ---------------------------------------------------------------------------
# Color and formatting
# ---------------------------------------------------------------------------

setup_colors() {
    if [[ -t 1 ]] && [[ "${NO_COLOR:-}" != "1" ]] && [[ "${TERM:-}" != "dumb" ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        CYAN='\033[0;36m'
        BOLD='\033[1m'
        DIM='\033[2m'
        RESET='\033[0m'
    else
        RED='' GREEN='' YELLOW='' BLUE='' CYAN='' BOLD='' DIM='' RESET=''
    fi
}

info()    { [[ "$QUIET" == true ]] && return; echo -e "${BLUE}[info]${RESET}    $*"; }
success() { [[ "$QUIET" == true ]] && return; echo -e "${GREEN}[ok]${RESET}      $*"; }
warn()    { echo -e "${YELLOW}[warn]${RESET}    $*" >&2; }
error()   { echo -e "${RED}[error]${RESET}   $*" >&2; }
step()    { [[ "$QUIET" == true ]] && return; echo -e "${CYAN}[step]${RESET}    ${BOLD}$*${RESET}"; }
detail()  { [[ "$VERBOSE" != true ]] && return; echo -e "${DIM}          $*${RESET}"; }

# ---------------------------------------------------------------------------
# Help
# ---------------------------------------------------------------------------

usage() {
    cat <<EOF
${BOLD}${SCRIPT_NAME}${RESET} v${VERSION} - Inspect IncusOS node health

${BOLD}USAGE${RESET}
    ${SCRIPT_NAME} [OPTIONS] ACTION

${BOLD}ACTIONS${RESET}
    -s, --status        Quick status overview (hostname, version, uptime,
                        memory, TPM, Secure Boot, trusted state, update)
    -p, --partitions    Storage and partition info (drives, ZFS pools,
                        volumes, encryption status)
    -t, --tpm           TPM and security details (TPM status, Secure Boot,
                        certificates, recovery keys, trusted state)
    -S, --services      Service configuration (iscsi, lvm, multipath, nvme,
                        ovn, tailscale, usbip; OVN details)
    -N, --network       Network configuration (interfaces, addresses, routes,
                        state, DNS)
    -u, --update        Update configuration and state (channel, frequency,
                        auto reboot, versions, last check)
    -a, --all           Run all of the above

${BOLD}OPTIONS${RESET}
    -r, --remote NAME   Incus remote (default: ${REMOTE})
    --node NAME         Specific node (default: all nodes in cluster)
    -j, --json          Output raw JSON instead of formatted text
    -n, --dry-run       Show what would be queried without executing
    -v, --verbose       Show detailed output
    -q, --quiet         Suppress informational output
    -h, --help          Show this help message

${BOLD}API ENDPOINTS${RESET}
    /os/1.0                     Environment (hostname, version, uptime)
    /os/1.0/system/security     TPM, Secure Boot, encryption
    /os/1.0/system/storage      Drives, ZFS pools, volumes
    /os/1.0/system/resources    CPU, memory
    /os/1.0/system/network      Interfaces, DNS
    /os/1.0/system/update       Update config and state
    /os/1.0/services/SERVICE    Per-service configuration

${BOLD}EXAMPLES${RESET}
    # Quick status of all cluster nodes
    ${SCRIPT_NAME} --status

    # Storage info for a specific node
    ${SCRIPT_NAME} --partitions --node oc-node-02

    # All info as JSON
    ${SCRIPT_NAME} --all --json

    # Different remote
    ${SCRIPT_NAME} --status --remote lab-node-01

    # Preview queries without running
    ${SCRIPT_NAME} --all --dry-run
EOF
}

# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -s|--status)      ACTION="status";;
            -p|--partitions)  ACTION="partitions";;
            -t|--tpm)         ACTION="tpm";;
            -S|--services)    ACTION="services";;
            -N|--network)     ACTION="network";;
            -u|--update)      ACTION="update";;
            -a|--all)         ACTION="all";;
            -r|--remote)      REMOTE="$2"; shift;;
            --node)           NODE="$2"; shift;;
            -j|--json)        JSON_OUTPUT=true;;
            -n|--dry-run)     DRY_RUN=true;;
            -v|--verbose)     VERBOSE=true;;
            -q|--quiet)       QUIET=true;;
            -h|--help)        usage; exit 0;;
            *)
                error "Unknown option: $1"
                echo "Run '${SCRIPT_NAME} --help' for usage." >&2
                exit 1
                ;;
        esac
        shift
    done

    if [[ -z "$ACTION" ]]; then
        error "No action specified"
        echo "Run '${SCRIPT_NAME} --help' for usage." >&2
        exit 1
    fi
}

# ---------------------------------------------------------------------------
# Node discovery
# ---------------------------------------------------------------------------

discover_nodes() {
    if [[ -n "$NODE" ]]; then
        echo "$NODE"
        return
    fi

    local members
    members=$(incus cluster list "${REMOTE}:" -f csv 2>/dev/null \
        | cut -d',' -f1 \
        | sed 's/^ *//;s/ *$//') || true

    if [[ -z "$members" ]]; then
        # Not a cluster or cannot list -- fall back to the remote itself
        echo "$REMOTE"
        return
    fi

    echo "$members"
}

# ---------------------------------------------------------------------------
# Query helper
# ---------------------------------------------------------------------------

query_node() {
    local node="$1" endpoint="$2"

    if [[ "$DRY_RUN" == true ]]; then
        info "(dry-run) incus query ${node}:${endpoint}"
        return 0
    fi

    detail "incus query ${node}:${endpoint}"
    incus query "${node}:${endpoint}" 2>/dev/null || echo "{}"
}

# ---------------------------------------------------------------------------
# Section: status
# ---------------------------------------------------------------------------

show_status() {
    local node="$1"

    local env_json resources_json security_json update_json
    env_json=$(query_node "$node" "/os/1.0")
    resources_json=$(query_node "$node" "/os/1.0/system/resources")
    security_json=$(query_node "$node" "/os/1.0/system/security")
    update_json=$(query_node "$node" "/os/1.0/system/update")

    if [[ "$DRY_RUN" == true ]]; then return; fi

    if [[ "$JSON_OUTPUT" == true ]]; then
        python3 -c "
import json, sys
print(json.dumps({
    'environment': json.loads(sys.argv[1]),
    'resources': json.loads(sys.argv[2]),
    'security': json.loads(sys.argv[3]),
    'update': json.loads(sys.argv[4]),
}, indent=2))
" "$env_json" "$resources_json" "$security_json" "$update_json"
        return
    fi

    python3 -c "
import json, sys

env = json.loads(sys.argv[1]).get('environment', json.loads(sys.argv[1]))
res = json.loads(sys.argv[2])
sec_raw = json.loads(sys.argv[3])
sec = sec_raw.get('state', sec_raw)
upd = json.loads(sys.argv[4])

hostname = env.get('hostname', 'unknown')
version = env.get('os_version', env.get('version', 'unknown'))
uptime_s = env.get('uptime', 0)

# Format uptime
days = uptime_s // 86400
hours = (uptime_s % 86400) // 3600
mins = (uptime_s % 3600) // 60
parts = []
if days: parts.append(f'{days}d')
if hours: parts.append(f'{hours}h')
parts.append(f'{mins}m')
uptime_str = ' '.join(parts)

# Memory
mem = res.get('memory', {})
mem_used = mem.get('used', 0)
mem_total = mem.get('total', 0)

def fmt_bytes(b):
    for unit in ['B','KiB','MiB','GiB','TiB']:
        if b < 1024:
            return f'{b:.1f} {unit}'
        b /= 1024
    return f'{b:.1f} PiB'

mem_pct = (mem_used / mem_total * 100) if mem_total else 0

# Security (state level)
tpm_status = sec.get('tpm_status', 'unknown')
secboot_status = 'enabled' if sec.get('secure_boot_enabled', False) else 'disabled'
trusted = sec.get('system_state_is_trusted', 'unknown')

# Update
needs_reboot = upd.get('state', {}).get('needs_reboot', False)
last_check = upd.get('state', {}).get('last_check', 'never')

print(f'  Hostname:       {hostname}')
print(f'  Version:        {version}')
print(f'  Uptime:         {uptime_str}')
print(f'  Memory:         {fmt_bytes(mem_used)} / {fmt_bytes(mem_total)} ({mem_pct:.0f}%)')
print(f'  TPM:            {tpm_status}')
print(f'  Secure Boot:    {secboot_status}')
print(f'  System trusted: {trusted}')
print(f'  Needs reboot:   {\"yes\" if needs_reboot else \"no\"}')
print(f'  Last update:    {last_check}')
" "$env_json" "$resources_json" "$security_json" "$update_json"
}

# ---------------------------------------------------------------------------
# Section: partitions
# ---------------------------------------------------------------------------

show_partitions() {
    local node="$1"

    local storage_json security_json
    storage_json=$(query_node "$node" "/os/1.0/system/storage")
    security_json=$(query_node "$node" "/os/1.0/system/security")

    if [[ "$DRY_RUN" == true ]]; then return; fi

    if [[ "$JSON_OUTPUT" == true ]]; then
        python3 -c "
import json, sys
print(json.dumps({
    'storage': json.loads(sys.argv[1]),
    'security': json.loads(sys.argv[2]),
}, indent=2))
" "$storage_json" "$security_json"
        return
    fi

    python3 -c "
import json, sys

stor_raw = json.loads(sys.argv[1])
sec_raw = json.loads(sys.argv[2])
stor = stor_raw.get('state', stor_raw)
sec = sec_raw.get('state', sec_raw)

def fmt_bytes(b):
    for unit in ['B','KiB','MiB','GiB','TiB']:
        if b < 1024:
            return f'{b:.1f} {unit}'
        b /= 1024
    return f'{b:.1f} PiB'

# Drives
disks = stor.get('drives', stor.get('disks', []))
if disks:
    print('  Drives:')
    for d in disks:
        model = d.get('model_name', d.get('model', 'unknown'))
        bus = d.get('bus', 'unknown')
        size = d.get('capacity_in_bytes', d.get('size', 0))
        boot = ' [boot]' if d.get('boot') else ''
        did = d.get('id', '?').split('/')[-1]
        pool = d.get('member_pool', '')
        pool_str = f' pool={pool}' if pool else ''
        print(f'    {did}: {model} ({bus}, {fmt_bytes(size)}){boot}{pool_str}')
else:
    print('  Drives: none found')

# ZFS pools
pools = stor.get('pools', [])
if pools:
    print('  ZFS pools:')
    for p in pools:
        name = p.get('name', 'unknown')
        state = p.get('state', 'unknown')
        ptype = p.get('type', 'unknown')
        alloc = p.get('pool_allocated_space_in_bytes', 0)
        total = p.get('raw_pool_size_in_bytes', p.get('usable_pool_size_in_bytes', 0))
        pct = (alloc / total * 100) if total else 0
        enc_key = p.get('encryption_key_status', 'N/A')
        print(f'    {name}: {state} ({ptype}, {fmt_bytes(alloc)} / {fmt_bytes(total)}, {pct:.0f}%) enc_key={enc_key}')
        for v in p.get('volumes', []):
            vname = v.get('name', '?')
            vuse = v.get('use', '?')
            vusage = v.get('usage_in_bytes', 0)
            vquota = v.get('quota_in_bytes', 0)
            quota_str = fmt_bytes(vquota) if vquota else 'no quota'
            print(f'      vol {vname} ({vuse}): {fmt_bytes(vusage)} / {quota_str}')
else:
    print('  ZFS pools: none found')

# Encryption
encrypted_vols = sec.get('encrypted_volumes', [])
if encrypted_vols:
    print('  Encryption:')
    for v in encrypted_vols:
        vol = v.get('volume', '?')
        state = v.get('state', '?')
        print(f'    {vol}: {state}')
else:
    print('  Encryption: not configured')
" "$storage_json" "$security_json"
}

# ---------------------------------------------------------------------------
# Section: tpm
# ---------------------------------------------------------------------------

show_tpm() {
    local node="$1"

    local security_json
    security_json=$(query_node "$node" "/os/1.0/system/security")

    if [[ "$DRY_RUN" == true ]]; then return; fi

    if [[ "$JSON_OUTPUT" == true ]]; then
        python3 -c "
import json, sys
print(json.dumps(json.loads(sys.argv[1]), indent=2))
" "$security_json"
        return
    fi

    python3 -c "
import json, sys

sec_raw = json.loads(sys.argv[1])
sec = sec_raw.get('state', sec_raw)
sec_cfg = sec_raw.get('config', {})

# TPM
tpm_status = sec.get('tpm_status', 'unknown')
print(f'  TPM status:     {tpm_status}')

# Secure Boot
sb_enabled = sec.get('secure_boot_enabled', False)
print(f'  Secure Boot:    {\"enabled\" if sb_enabled else \"disabled\"}')

# Certificates
certs = sec.get('secure_boot_certificates', [])
if certs:
    print('  Certificates:')
    for c in certs:
        ctype = c.get('type', '?')
        subject = c.get('subject', 'unknown')
        fp = c.get('fingerprint', '')
        fp_short = fp[:16] + '...' if len(fp) > 16 else fp
        print(f'    {ctype:4s} {subject}')
        if fp_short:
            print(f'         fingerprint: {fp_short}')

# Encrypted volumes
enc_vols = sec.get('encrypted_volumes', [])
if enc_vols:
    print('  Encrypted volumes:')
    for v in enc_vols:
        print(f'    {v.get(\"volume\", \"?\")}: {v.get(\"state\", \"?\")}')

# Recovery keys
recovery_keys = sec_cfg.get('encryption_recovery_keys', [])
retrieved = sec.get('encryption_recovery_keys_retrieved', False)
print(f'  Recovery keys:  {len(recovery_keys)} key(s), retrieved={retrieved}')

# Pool recovery keys
pool_keys = sec.get('pool_recovery_keys', {})
if pool_keys:
    print(f'  Pool keys:      {list(pool_keys.keys())}')

# System trusted
trusted = sec.get('system_state_is_trusted', 'unknown')
print(f'  System trusted: {trusted}')
" "$security_json"
}

# ---------------------------------------------------------------------------
# Section: services
# ---------------------------------------------------------------------------

show_services() {
    local node="$1"

    if [[ "$JSON_OUTPUT" == true ]]; then
        local all_json="{"
        local first=true
        for svc in "${SERVICES[@]}"; do
            local svc_json
            svc_json=$(query_node "$node" "/os/1.0/services/${svc}")
            if [[ "$DRY_RUN" == true ]]; then continue; fi
            if [[ "$first" == true ]]; then
                first=false
            else
                all_json+=","
            fi
            all_json+="\"${svc}\":${svc_json}"
        done
        all_json+="}"
        if [[ "$DRY_RUN" != true ]]; then
            echo "$all_json" | python3 -m json.tool
        fi
        return
    fi

    for svc in "${SERVICES[@]}"; do
        local svc_json
        svc_json=$(query_node "$node" "/os/1.0/services/${svc}")
        if [[ "$DRY_RUN" == true ]]; then continue; fi

        python3 -c "
import json, sys

svc_name = sys.argv[1]
data = json.loads(sys.argv[2])

config = data.get('config', data)
enabled = config.get('enabled', False)
status = 'enabled' if enabled else 'disabled'

line = f'  {svc_name:12s} {status}'
print(line)

# Extra details for OVN
if svc_name == 'ovn' and enabled:
    db = config.get('database', '')
    tunnel_addr = config.get('tunnel_address', '')
    tunnel_proto = config.get('tunnel_protocol', '')
    if db:
        print(f'               database:        {db}')
    if tunnel_addr:
        print(f'               tunnel address:  {tunnel_addr}')
    if tunnel_proto:
        print(f'               tunnel protocol: {tunnel_proto}')
" "$svc" "$svc_json"
    done
}

# ---------------------------------------------------------------------------
# Section: network
# ---------------------------------------------------------------------------

show_network() {
    local node="$1"

    local network_json
    network_json=$(query_node "$node" "/os/1.0/system/network")

    if [[ "$DRY_RUN" == true ]]; then return; fi

    if [[ "$JSON_OUTPUT" == true ]]; then
        echo "$network_json" | python3 -m json.tool
        return
    fi

    python3 -c "
import json, sys

net = json.loads(sys.argv[1])

def fmt_bytes(b):
    for unit in ['B','KiB','MiB','GiB','TiB']:
        if b < 1024: return f'{b:.1f} {unit}'
        b /= 1024
    return f'{b:.1f} PiB'

# Interfaces
ifaces = net.get('interfaces', net.get('config', {}).get('interfaces', []))
if not isinstance(ifaces, list): ifaces = []
print('  Interfaces:')
for iface in ifaces:
    name = iface.get('name', 'unknown')
    hwaddr = iface.get('hwaddr', iface.get('mac', ''))
    roles = iface.get('roles', [])
    roles_str = ', '.join(roles) if isinstance(roles, list) and roles else ''
    addrs = iface.get('addresses', [])
    addr_strs = [a if isinstance(a, str) else a.get('address', '?') for a in addrs]
    state = iface.get('state', {})

    header = f'    {name}'
    if hwaddr: header += f' ({hwaddr})'
    if roles_str: header += f'  [{roles_str}]'
    print(header)

    if addr_strs:
        print(f'      addresses: {\", \".join(addr_strs)}')

    link_state = state.get('state', state.get('link', ''))
    mtu = state.get('mtu', iface.get('mtu', ''))
    speed = state.get('speed', '')
    if link_state:
        parts = [f'state={link_state}']
        if mtu: parts.append(f'mtu={mtu}')
        if speed: parts.append(f'speed={speed}')
        print(f'      link:      {\", \".join(parts)}')

    routes = iface.get('routes', [])
    if isinstance(routes, list) and routes:
        route_strs = []
        for r in routes:
            if isinstance(r, str): route_strs.append(r)
            elif isinstance(r, dict):
                s = r.get('destination', r.get('to', '?'))
                via = r.get('via', r.get('gateway', ''))
                if via: s += f' via {via}'
                route_strs.append(s)
        if route_strs:
            print(f'      routes:    {route_strs[0]}')
            for rs in route_strs[1:]:
                print(f'                 {rs}')

    stats = state.get('stats', iface.get('stats', {}))
    if isinstance(stats, dict) and stats:
        rx = stats.get('rx_bytes', stats.get('rx', 0))
        tx = stats.get('tx_bytes', stats.get('tx', 0))
        if rx or tx:
            print(f'      stats:     rx={fmt_bytes(rx)}, tx={fmt_bytes(tx)}')

# DNS
dns = net.get('dns', net.get('config', {}).get('dns', {}))
if isinstance(dns, dict) and dns:
    print('  DNS:')
    hostname = dns.get('hostname', dns.get('search', ''))
    nameservers = dns.get('nameservers', dns.get('servers', []))
    if hostname: print(f'    hostname:    {hostname}')
    if isinstance(nameservers, list) and nameservers:
        print(f'    nameservers: {\", \".join(str(n) for n in nameservers)}')
" "$network_json"
}

# ---------------------------------------------------------------------------
# Section: update
# ---------------------------------------------------------------------------

show_update() {
    local node="$1"

    local update_json env_json
    update_json=$(query_node "$node" "/os/1.0/system/update")
    env_json=$(query_node "$node" "/os/1.0")

    if [[ "$DRY_RUN" == true ]]; then return; fi

    if [[ "$JSON_OUTPUT" == true ]]; then
        echo "$update_json" | python3 -m json.tool
        return
    fi

    python3 -c "
import json, sys

upd = json.loads(sys.argv[1])
env = json.loads(sys.argv[2]).get('environment', json.loads(sys.argv[2]))
config = upd.get('config', {})
state = upd.get('state', {})

channel = config.get('channel', 'unknown')
frequency = config.get('check_frequency', 'unknown')
auto_reboot = config.get('auto_reboot', 'unknown')

current = env.get('os_version', 'unknown')
next_ver = env.get('os_version_next', 'none')
update_status = state.get('status', '')
last_check = state.get('last_check', 'never')
needs_reboot = state.get('needs_reboot', False)

print(f'  Channel:        {channel}')
print(f'  Check freq:     {frequency}')
print(f'  Auto reboot:    {auto_reboot}')
print(f'  Current:        {current}')
print(f'  Next version:   {next_ver}')
print(f'  Update status:  {update_status}')
print(f'  Last check:     {last_check}')
print(f'  Needs reboot:   {\"yes\" if needs_reboot else \"no\"}')
" "$update_json" "$env_json"
}

# ---------------------------------------------------------------------------
# Run a section across all nodes
# ---------------------------------------------------------------------------

run_section() {
    local section_name="$1"
    local section_fn="$2"
    shift 2
    local nodes=("$@")

    step "$section_name"

    for node in "${nodes[@]}"; do
        echo -e "  ${BOLD}--- ${node} ---${RESET}"
        "$section_fn" "$node"
        echo
    done
}

run_section_json() {
    local section_fn="$1"
    shift
    local nodes=("$@")

    if [[ ${#nodes[@]} -eq 1 ]]; then
        "$section_fn" "${nodes[0]}"
    else
        echo "{"
        local first=true
        for node in "${nodes[@]}"; do
            if [[ "$first" == true ]]; then
                first=false
            else
                echo ","
            fi
            echo "  \"${node}\":"
            "$section_fn" "$node" | sed 's/^/  /'
        done
        echo "}"
    fi
}

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

main() {
    setup_colors
    parse_args "$@"

    # Discover cluster nodes
    if [[ "$DRY_RUN" == true ]]; then
        if [[ -n "$NODE" ]]; then
            info "(dry-run) Would query node: ${NODE}"
        else
            info "(dry-run) Would discover nodes via: incus cluster list ${REMOTE}: -f csv"
        fi
    fi

    local nodes_raw
    nodes_raw=$(discover_nodes)

    local nodes=()
    while IFS= read -r line; do
        [[ -n "$line" ]] && nodes+=("$line")
    done <<< "$nodes_raw"

    if [[ ${#nodes[@]} -eq 0 ]]; then
        error "No nodes found. Check that remote '${REMOTE}' is accessible."
        error "  incus remote list"
        exit 1
    fi

    if [[ "$DRY_RUN" != true ]]; then
        info "Nodes: ${nodes[*]}"
    fi
    echo

    # Section lookup: action -> (title, function)
    local -A SECTION_TITLE=(
        [status]="Status" [partitions]="Storage & Partitions"
        [tpm]="TPM & Security" [services]="Services"
        [network]="Network" [update]="Update"
    )

    # Build section list
    local sections=()
    if [[ "$ACTION" == "all" ]]; then
        sections=(status partitions tpm services network update)
    else
        sections=("$ACTION")
    fi

    # Dispatch
    if [[ "$JSON_OUTPUT" == true && ${#sections[@]} -gt 1 ]]; then
        echo "{"
        local i=0
        for section in "${sections[@]}"; do
            echo "  \"${section}\":"
            run_section_json "show_${section}" "${nodes[@]}" | sed 's/^/  /'
            i=$((i + 1))
            [[ $i -lt ${#sections[@]} ]] && echo ","
        done
        echo "}"
    else
        for section in "${sections[@]}"; do
            if [[ "$JSON_OUTPUT" == true ]]; then
                run_section_json "show_${section}" "${nodes[@]}"
            else
                run_section "${SECTION_TITLE[$section]}" "show_${section}" "${nodes[@]}"
            fi
        done
    fi
}

main "$@"