#!/usr/bin/env bash
set -euo pipefail

readonly VERSION="0.1.0"
readonly SCRIPT_NAME="lab-status"
readonly SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
readonly REPO_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
readonly ENVS_DIR="${REPO_DIR}/envs"

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

QUIET=false
VERBOSE=false
ENV_NAME=""
NO_CHECK=false
JSON_OUTPUT=false

# Service arrays (parallel indexed)
declare -a SVC_NAMES=()
declare -a SVC_URLS=()
declare -a SVC_STATUSES=()

# --- Colors & logging --------------------------------------------------------

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; }
detail()  { [[ "$VERBOSE" != true ]] && return; echo -e "${DIM}          $*${RESET}"; }

# --- Usage --------------------------------------------------------------------

usage() {
    cat <<EOF
${BOLD}${SCRIPT_NAME}${RESET} v${VERSION} - Lab service status overview

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

${BOLD}OPTIONS${RESET}
    -e, --env NAME        Environment name (auto-detects if only one exists)
        --no-check        Skip reachability checks, just show URLs
        --json            Machine-readable JSON output
    -v, --verbose         Show extra details
    -q, --quiet           Suppress info messages
    -h, --help            Show this help

${BOLD}DESCRIPTION${RESET}
    Reads YAML configs from envs/<env>/ and checks whether each service
    is reachable. Any HTTP response (even 401/403) counts as UP; connection
    timeout or refused counts as DOWN.

    Services detected: Proxmox, OC Server, Cluster Nodes, Aether, AWX,
    HAProxy VIP, Grafana, Prometheus.

${BOLD}EXAMPLES${RESET}
    # Show hetzner lab with reachability checks
    ${SCRIPT_NAME} --env hetzner

    # Show beelink lab without checks
    ${SCRIPT_NAME} --env beelink --no-check

    # JSON output for scripting
    ${SCRIPT_NAME} --env hetzner --json

EOF
}

# --- YAML helpers -------------------------------------------------------------

# Extract a value from a YAML section: yaml_val FILE SECTION KEY
yaml_val() {
    local file="$1" section="$2" key="$3"
    grep -A50 "^${section}:" "$file" \
        | grep "^ *${key}:" | head -1 \
        | sed "s/^ *${key}: *//" | tr -d '"' | sed 's/ *#.*//' || true
}

# Extract a top-level YAML value: yaml_val_top FILE KEY
yaml_val_top() {
    local file="$1" key="$2"
    grep "^${key}:" "$file" | head -1 \
        | sed "s/^${key}: *//" | tr -d '"' | sed 's/ *#.*//' || true
}

# --- Service collection -------------------------------------------------------

add_service() {
    local name="$1" url="$2"
    SVC_NAMES+=("$name")
    SVC_URLS+=("$url")
    SVC_STATUSES+=("")
}

detect_env() {
    if [[ -n "$ENV_NAME" ]]; then
        if [[ ! -d "${ENVS_DIR}/${ENV_NAME}" ]]; then
            error "Environment not found: ${ENV_NAME}"
            local available
            available=$(ls "$ENVS_DIR" 2>/dev/null | tr '\n' ' ')
            [[ -n "$available" ]] && error "Available: ${available}"
            exit 1
        fi
        return
    fi

    local envs=()
    for d in "$ENVS_DIR"/*/; do
        [[ -d "$d" ]] && envs+=("$(basename "$d")")
    done

    if [[ ${#envs[@]} -eq 1 ]]; then
        ENV_NAME="${envs[0]}"
        info "Auto-detected environment: ${ENV_NAME}"
    elif [[ ${#envs[@]} -eq 0 ]]; then
        error "No environments found in ${ENVS_DIR}/"
        exit 1
    else
        error "Multiple environments found: ${envs[*]}"
        error "Use --env NAME to select one"
        exit 1
    fi
}

collect_proxmox() {
    local env_dir="${ENVS_DIR}/${ENV_NAME}"
    local pve_file=""

    if [[ -f "${env_dir}/proxmox.yaml" ]]; then
        pve_file="${env_dir}/proxmox.yaml"
    elif [[ -f "${env_dir}/proxmox.yaml.example" ]]; then
        pve_file="${env_dir}/proxmox.yaml.example"
        warn "Using proxmox.yaml.example (proxmox.yaml not found)"
    else
        detail "No proxmox.yaml found, skipping Proxmox"
        return
    fi

    local host
    host=$(yaml_val_top "$pve_file" "host")
    if [[ -n "$host" ]]; then
        host="${host%%/*}"
        add_service "Proxmox" "https://${host}:8006"
        detail "Proxmox from $(basename "$pve_file"): ${host}"
    fi
}

collect_lab_vms() {
    local env_dir="${ENVS_DIR}/${ENV_NAME}"
    local lab_file=""

    # Prefer lab-production.yaml over lab-cluster.yaml
    if [[ -f "${env_dir}/lab-production.yaml" ]]; then
        lab_file="${env_dir}/lab-production.yaml"
    elif [[ -f "${env_dir}/lab-cluster.yaml" ]]; then
        lab_file="${env_dir}/lab-cluster.yaml"
    else
        detail "No lab-*.yaml found, skipping cluster VMs"
        return
    fi

    detail "Lab VMs from $(basename "$lab_file")"

    local vm_data
    vm_data=$(python3 -c "
import sys, re

try:
    import yaml
    with open(sys.argv[1]) as f:
        data = yaml.safe_load(f)
except ImportError:
    data = {'vms': []}
    with open(sys.argv[1]) as f:
        content = f.read()
    in_vms = False
    current = {}
    for line in content.split('\n'):
        if line.startswith('vms:'):
            in_vms = True
            continue
        if in_vms:
            if line.startswith('  - '):
                if current:
                    data['vms'].append(current)
                current = {}
                line = line[4:]
            elif line.startswith('    '):
                line = line[4:]
            elif line and not line.startswith(' ') and not line.startswith('#'):
                if current:
                    data['vms'].append(current)
                break
            else:
                continue
            m = re.match(r'(\w+):\s*(.*)', line.strip())
            if m:
                current[m.group(1)] = m.group(2).strip('\"')
    if current:
        data['vms'].append(current)

for vm in data.get('vms', []):
    name = vm.get('name', '')
    app = vm.get('app', '')
    ip = vm.get('ip', '').split('/')[0]
    print(f'{name}|{app}|{ip}')
" "$lab_file" 2>/dev/null) || return

    local node_idx=0
    while IFS='|' read -r name app ip; do
        [[ -z "$ip" ]] && continue
        case "$app" in
            operations-center)
                add_service "OC Server" "https://${ip}:8443"
                ;;
            incus)
                node_idx=$((node_idx + 1))
                local label
                label=$(printf "Cluster Node %02d" "$node_idx")
                add_service "$label" "https://${ip}:8443"
                ;;
        esac
    done <<< "$vm_data"
}

collect_haproxy() {
    local env_dir="${ENVS_DIR}/${ENV_NAME}"
    local file="${env_dir}/haproxy.yaml"
    [[ -f "$file" ]] || return

    local aether_url vip
    aether_url=$(yaml_val "$file" "haproxy" "aether_url")
    vip=$(yaml_val "$file" "haproxy" "vip")

    if [[ -n "$aether_url" ]]; then
        add_service "Aether" "$aether_url"
        detail "Aether from haproxy.yaml: ${aether_url}"
    fi
    if [[ -n "$vip" ]]; then
        vip="${vip%%/*}"
        add_service "HAProxy VIP" "http://${vip}:80"
        detail "HAProxy VIP from haproxy.yaml: ${vip}"
    fi
}

collect_awx() {
    local env_dir="${ENVS_DIR}/${ENV_NAME}"
    local file="${env_dir}/awx.yaml"
    [[ -f "$file" ]] || return

    local ip
    ip=$(yaml_val "$file" "awx" "ip")
    [[ -z "$ip" ]] && return
    ip="${ip%%/*}"
    add_service "AWX" "http://${ip}:30080"
    detail "AWX from awx.yaml: ${ip}"
}

collect_observability() {
    local env_dir="${ENVS_DIR}/${ENV_NAME}"
    local file="${env_dir}/observability.yaml"
    [[ -f "$file" ]] || return

    local forward_vip
    forward_vip=$(yaml_val "$file" "observability" "forward_vip")
    [[ -z "$forward_vip" ]] && return
    forward_vip="${forward_vip%%/*}"

    add_service "Grafana" "http://${forward_vip}:3000"
    add_service "Prometheus" "http://${forward_vip}:9090"
    detail "Observability from observability.yaml: forward_vip=${forward_vip}"
}

# --- Reachability checks -----------------------------------------------------

check_reachability() {
    local url="$1" result_file="$2"
    local http_code
    http_code=$(curl -sk --connect-timeout 2 --max-time 3 \
        -o /dev/null -w '%{http_code}' "$url" 2>/dev/null) || http_code="000"
    if [[ "$http_code" != "000" ]]; then
        echo "UP" > "$result_file"
    else
        echo "DOWN" > "$result_file"
    fi
}

run_checks() {
    local tmpdir
    tmpdir=$(mktemp -d "${TMPDIR:-/tmp}/${SCRIPT_NAME}-XXXXXX")
    # shellcheck disable=SC2064
    trap "rm -rf '$tmpdir'" EXIT

    local i
    for i in "${!SVC_URLS[@]}"; do
        check_reachability "${SVC_URLS[$i]}" "${tmpdir}/${i}" &
    done
    wait

    for i in "${!SVC_URLS[@]}"; do
        if [[ -f "${tmpdir}/${i}" ]]; then
            SVC_STATUSES[$i]=$(cat "${tmpdir}/${i}")
        else
            SVC_STATUSES[$i]="DOWN"
        fi
    done
}

# --- Output -------------------------------------------------------------------

print_table() {
    local max_name=7 max_url=3  # minimum: "Service" / "URL"
    local i
    for i in "${!SVC_NAMES[@]}"; do
        [[ ${#SVC_NAMES[$i]} -gt $max_name ]] && max_name=${#SVC_NAMES[$i]}
        [[ ${#SVC_URLS[$i]} -gt $max_url ]] && max_url=${#SVC_URLS[$i]}
    done

    local total_width=$((max_name + max_url + 14))
    local sep
    sep=$(printf '%*s' "$total_width" '' | tr ' ' '─')

    echo
    echo -e "${BOLD}Lab Environment: ${ENV_NAME}${RESET}"
    echo "$sep"
    printf "%-${max_name}s  %-${max_url}s  %s\n" "Service" "URL" "Status"
    echo "$sep"

    for i in "${!SVC_NAMES[@]}"; do
        local status="${SVC_STATUSES[$i]:-—}"
        local status_colored
        case "$status" in
            UP)   status_colored="${GREEN}UP${RESET}" ;;
            DOWN) status_colored="${RED}DOWN${RESET}" ;;
            *)    status_colored="$status" ;;
        esac
        printf "%-${max_name}s  %-${max_url}s  " "${SVC_NAMES[$i]}" "${SVC_URLS[$i]}"
        echo -e "$status_colored"
    done

    echo "$sep"
    echo
}

print_json() {
    local last=$((${#SVC_NAMES[@]} - 1))
    echo "["
    local i
    for i in "${!SVC_NAMES[@]}"; do
        local comma=","
        [[ $i -eq $last ]] && comma=""
        local status_json
        if [[ -n "${SVC_STATUSES[$i]:-}" ]]; then
            status_json="\"${SVC_STATUSES[$i]}\""
        else
            status_json="null"
        fi
        printf '  {"service": "%s", "url": "%s", "status": %s}%s\n' \
            "${SVC_NAMES[$i]}" "${SVC_URLS[$i]}" "$status_json" "$comma"
    done
    echo "]"
}

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

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -e|--env)
                [[ $# -lt 2 ]] && { error "--env requires a value"; exit 1; }
                ENV_NAME="$2"; shift
                ;;
            --no-check)    NO_CHECK=true ;;
            --json)        JSON_OUTPUT=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."
                exit 1
                ;;
        esac
        shift
    done
}

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

main() {
    setup_colors
    parse_args "$@"
    detect_env

    # Collect services from all config files
    collect_proxmox
    collect_lab_vms
    collect_haproxy
    collect_awx
    collect_observability

    if [[ ${#SVC_NAMES[@]} -eq 0 ]]; then
        warn "No services found in envs/${ENV_NAME}/"
        exit 0
    fi

    # Run reachability checks (parallel)
    if [[ "$NO_CHECK" != true ]]; then
        info "Checking ${#SVC_NAMES[@]} services..."
        run_checks
    fi

    # Output
    if [[ "$JSON_OUTPUT" == true ]]; then
        print_json
    else
        print_table
    fi
}

main "$@"