#!/bin/bash
# Service Health Check Script for Edge Sensor
#
# Verifies all critical services are running properly after installation/reboot.
# Run manually with: sudo ./check-services.sh
#
# Exit codes:
#   0 - All services healthy
#   1 - One or more critical services failed
#   2 - One or more optional services failed (non-fatal)

set -euo pipefail

# ============================================================================
# Color output
# ============================================================================

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# ============================================================================
# Service Definitions
# ============================================================================

# Critical services - must be running for proper operation
CRITICAL_SERVICES=(
    "edge-sensor.service"
    "web-portal.service"
)

# Important services - should be running but may not be installed
IMPORTANT_SERVICES=(
    "power-monitor.service"
    "wittypi.service"
    "tailscaled.service"
)

# Optional services - 4G modem (may not be present on all units)
OPTIONAL_SERVICES=(
    "ModemManager.service"
    "NetworkManager.service"
)

# ============================================================================
# Helper Functions
# ============================================================================

print_header() {
    echo -e "${BLUE}========================================${NC}"
    echo -e "${BLUE}$1${NC}"
    echo -e "${BLUE}========================================${NC}"
    echo ""
}

print_service_status() {
    local service=$1
    local category=$2  # "critical", "important", "optional"

    # Check if service unit exists
    if ! systemctl list-unit-files "$service" &>/dev/null; then
        echo -e "${YELLOW}⚠  $service${NC} - Not installed"
        return 2
    fi

    # Check if service is active
    if systemctl is-active "$service" &>/dev/null; then
        echo -e "${GREEN}✓ $service${NC} - Running"

        # Check for recent errors in logs (last 50 lines)
        local error_count=$(journalctl -u "$service" -n 50 --no-pager 2>/dev/null | grep -iE "error|failed|fatal|critical" | wc -l || echo 0)

        if [ "$error_count" -gt 0 ]; then
            echo -e "  ${YELLOW}⚠  Found $error_count error(s) in recent logs${NC}"
            echo -e "  ${BLUE}→ Check with: sudo journalctl -u $service -n 100${NC}"
        fi

        return 0
    else
        # Service is not active
        local state=$(systemctl is-active "$service" 2>&1 || echo "unknown")

        if [ "$category" = "critical" ]; then
            echo -e "${RED}✗ $service${NC} - $state (CRITICAL)"
            echo -e "  ${BLUE}→ Check with: sudo systemctl status $service${NC}"
            return 1
        elif [ "$category" = "important" ]; then
            echo -e "${YELLOW}⚠  $service${NC} - $state"
            echo -e "  ${BLUE}→ Check with: sudo systemctl status $service${NC}"
            return 2
        else
            echo -e "${YELLOW}⚠  $service${NC} - $state (optional)"
            return 2
        fi
    fi
}

check_tailscale_connectivity() {
    echo ""
    echo -e "${BLUE}Checking Tailscale connectivity...${NC}"

    if ! command -v tailscale &>/dev/null; then
        echo -e "${YELLOW}⚠  tailscale CLI not found${NC}"
        return 2
    fi

    if tailscale status --json &>/dev/null; then
        local ip=$(tailscale ip -4 2>/dev/null | head -n1)
        local hostname=$(tailscale status --json 2>/dev/null | grep -o '"HostName":"[^"]*"' | cut -d'"' -f4 || echo "unknown")

        if [ -n "$ip" ]; then
            echo -e "${GREEN}✓ Tailscale connected${NC}"
            echo -e "  IP: $ip"
            echo -e "  Hostname: $hostname"
            return 0
        else
            echo -e "${YELLOW}⚠  Tailscale running but no IP assigned${NC}"
            return 2
        fi
    else
        echo -e "${RED}✗ Tailscale not connected${NC}"
        echo -e "  ${BLUE}→ Authenticate with: sudo tailscale up${NC}"
        return 1
    fi
}

check_wittypi_hardware() {
    echo ""
    echo -e "${BLUE}Checking WittyPi hardware...${NC}"

    if ! command -v i2cdetect &>/dev/null; then
        echo -e "${YELLOW}⚠  i2c-tools not installed${NC}"
        return 2
    fi

    # Check for WittyPi on I2C bus 1, address 0x08
    if i2cdetect -y 1 2>/dev/null | grep -qE '(^|[[:space:]])08([[:space:]]|$)'; then
        echo -e "${GREEN}✓ WittyPi detected on I2C bus 1 (0x08)${NC}"

        # Check for INA219 battery monitor on address 0x40
        if i2cdetect -y 1 2>/dev/null | grep -qE '(^|[[:space:]])40([[:space:]]|$)'; then
            echo -e "${GREEN}✓ INA219 battery monitor detected (0x40)${NC}"
        else
            echo -e "${YELLOW}⚠  INA219 battery monitor not detected (0x40)${NC}"
            echo -e "  ${BLUE}→ Battery voltage monitoring may not work${NC}"
        fi

        return 0
    else
        echo -e "${YELLOW}⚠  WittyPi not detected on I2C bus${NC}"
        echo -e "  ${BLUE}→ Check hardware connection${NC}"
        return 2
    fi
}

check_power_monitor_data() {
    echo ""
    echo -e "${BLUE}Checking power monitoring data...${NC}"

    local data_dir="/var/log/power_daily"

    if [ ! -d "$data_dir" ]; then
        echo -e "${YELLOW}⚠  Power data directory not found: $data_dir${NC}"
        return 2
    fi

    # Check for today's data file
    local today_file="$data_dir/power_$(date +%Y-%m-%d).csv"

    if [ -f "$today_file" ]; then
        local line_count=$(wc -l < "$today_file" 2>/dev/null || echo 0)
        local file_age_minutes=$(( ($(date +%s) - $(stat -c %Y "$today_file" 2>/dev/null || date +%s)) / 60 ))

        echo -e "${GREEN}✓ Today's power data file exists${NC}"
        echo -e "  File: $today_file"
        echo -e "  Lines: $line_count"
        echo -e "  Last updated: $file_age_minutes minutes ago"

        if [ "$file_age_minutes" -gt 5 ]; then
            echo -e "  ${YELLOW}⚠  File not updated recently (> 5 min)${NC}"
            echo -e "  ${BLUE}→ Check if power-monitor service is running${NC}"
        fi

        return 0
    else
        echo -e "${YELLOW}⚠  No power data for today${NC}"
        echo -e "  ${BLUE}→ Expected: $today_file${NC}"
        return 2
    fi
}

check_network_connectivity() {
    echo ""
    echo -e "${BLUE}Checking network connectivity...${NC}"

    # Check for any active network interfaces (excluding loopback)
    local active_interfaces=$(ip link show | grep -E "state UP" | grep -v "lo:" | wc -l)

    if [ "$active_interfaces" -gt 0 ]; then
        echo -e "${GREEN}✓ $active_interfaces active network interface(s)${NC}"

        # Try to ping Google DNS (8.8.8.8) to verify internet connectivity
        if ping -c 1 -W 2 8.8.8.8 &>/dev/null; then
            echo -e "${GREEN}✓ Internet connectivity verified${NC}"
            return 0
        else
            echo -e "${YELLOW}⚠  No internet connectivity${NC}"
            echo -e "  ${BLUE}→ Check WiFi/4G configuration${NC}"
            return 2
        fi
    else
        echo -e "${RED}✗ No active network interfaces${NC}"
        echo -e "  ${BLUE}→ Check network configuration${NC}"
        return 1
    fi
}

check_disk_space() {
    echo ""
    echo -e "${BLUE}Checking disk space...${NC}"

    local root_usage=$(df -h / | awk 'NR==2 {print $5}' | sed 's/%//')
    local root_avail=$(df -h / | awk 'NR==2 {print $4}')

    echo -e "  Root filesystem: ${root_usage}% used ($root_avail available)"

    if [ "$root_usage" -gt 90 ]; then
        echo -e "${RED}✗ Disk space critically low (<10% free)${NC}"
        return 1
    elif [ "$root_usage" -gt 80 ]; then
        echo -e "${YELLOW}⚠  Disk space running low (<20% free)${NC}"
        return 2
    else
        echo -e "${GREEN}✓ Disk space OK${NC}"
        return 0
    fi
}

check_system_load() {
    echo ""
    echo -e "${BLUE}Checking system load...${NC}"

    local load_1min=$(uptime | awk -F'load average:' '{print $2}' | awk '{print $1}' | sed 's/,//')
    local cpu_count=$(nproc)

    echo -e "  Load average (1min): $load_1min"
    echo -e "  CPU cores: $cpu_count"

    # Load average threshold: 2x CPU count (high but not critical for edge device)
    local load_threshold=$(echo "$cpu_count * 2" | bc 2>/dev/null || echo 8)

    if (( $(echo "$load_1min > $load_threshold" | bc -l 2>/dev/null || echo 0) )); then
        echo -e "${YELLOW}⚠  System load is high${NC}"
        return 2
    else
        echo -e "${GREEN}✓ System load OK${NC}"
        return 0
    fi
}

check_cpu_governor() {
    echo ""
    echo -e "${BLUE}Checking CPU power management...${NC}"

    # Check if cpufreq is available
    if [ ! -d /sys/devices/system/cpu/cpu0/cpufreq ]; then
        echo -e "${YELLOW}⚠  CPU frequency scaling not available${NC}"
        return 2
    fi

    local all_powersave=true
    local cpu_governors=""

    # Check all CPU cores
    for cpu_dir in /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor; do
        if [ -f "$cpu_dir" ]; then
            local governor=$(cat "$cpu_dir" 2>/dev/null || echo "unknown")
            local cpu_num=$(echo "$cpu_dir" | grep -oP 'cpu\K[0-9]+')

            if [ "$governor" != "powersave" ]; then
                all_powersave=false
                cpu_governors="${cpu_governors}CPU${cpu_num}:$governor "
            fi
        fi
    done

    if [ "$all_powersave" = true ]; then
        local cpu_count=$(nproc)
        echo -e "${GREEN}✓ All CPUs in powersave mode ($cpu_count cores)${NC}"
        return 0
    else
        echo -e "${RED}✗ CPUs not in powersave mode${NC}"
        echo -e "  ${cpu_governors}"
        echo -e "  ${BLUE}→ Run: sudo ./install.sh configure_power_management${NC}"
        return 1
    fi
}

check_wittypi_power_on_config() {
    echo ""
    echo -e "${BLUE}Checking WittyPi power-on configuration...${NC}"

    if ! command -v i2cget &>/dev/null; then
        echo -e "${YELLOW}⚠  i2c-tools not installed${NC}"
        return 2
    fi

    # Check if WittyPi is present on I2C bus (0x08)
    if ! i2cdetect -y 1 2>/dev/null | grep -qE '(^|[[:space:]])08([[:space:]]|$)'; then
        echo -e "${YELLOW}⚠  WittyPi not detected, skipping power-on check${NC}"
        return 2
    fi

    # Read register 0x18: Default state when powered
    # 0x01 = ON at power (desired)
    # 0x00 = OFF at power (requires button press)
    local power_on_state=$(i2cget -y 1 0x08 0x18 2>/dev/null || echo "0xff")

    case "$power_on_state" in
        "0x01")
            echo -e "${GREEN}✓ WittyPi configured for ON at power${NC}"
            echo -e "  Device will start automatically when power applied"
            return 0
            ;;
        "0x00")
            echo -e "${RED}✗ WittyPi configured for OFF at power${NC}"
            echo -e "  Device requires button press to start (not suitable for remote deployment)"
            echo -e "  ${BLUE}→ Reconfigure with: sudo ./install.sh configure_wittypi${NC}"
            return 1
            ;;
        *)
            echo -e "${YELLOW}⚠  Could not read WittyPi power-on configuration (register 0x18: $power_on_state)${NC}"
            echo -e "  ${BLUE}→ Verify I2C connection and permissions${NC}"
            return 2
            ;;
    esac
}

# ============================================================================
# Main Health Check
# ============================================================================

main() {
    local critical_failed=0
    local warnings=0

    print_header "Edge Sensor Service Health Check"

    echo -e "${BLUE}Hostname:${NC} $(hostname)"
    echo -e "${BLUE}Uptime:${NC}   $(uptime -p)"
    echo -e "${BLUE}Date:${NC}     $(date)"
    echo ""

    # Check critical services
    print_header "Critical Services"
    for service in "${CRITICAL_SERVICES[@]}"; do
        if ! print_service_status "$service" "critical"; then
            if [ $? -eq 1 ]; then
                ((critical_failed++))
            else
                ((warnings++))
            fi
        fi
    done
    echo ""

    # Check important services
    print_header "Important Services"
    for service in "${IMPORTANT_SERVICES[@]}"; do
        if ! print_service_status "$service" "important"; then
            ((warnings++))
        fi
    done

    # Check Tailscale connectivity
    if ! check_tailscale_connectivity; then
        if [ $? -eq 1 ]; then
            ((critical_failed++))
        else
            ((warnings++))
        fi
    fi

    # Check WittyPi hardware
    if ! check_wittypi_hardware; then
        ((warnings++))
    fi

    # Check power monitoring data
    if ! check_power_monitor_data; then
        ((warnings++))
    fi
    echo ""

    # Check optional services
    print_header "Optional Services (4G Modem)"
    for service in "${OPTIONAL_SERVICES[@]}"; do
        if ! print_service_status "$service" "optional"; then
            : # Don't count optional service failures
        fi
    done
    echo ""

    # System health checks
    print_header "System Health"

    if ! check_cpu_governor; then
        if [ $? -eq 1 ]; then
            ((critical_failed++))
        else
            ((warnings++))
        fi
    fi

    if ! check_network_connectivity; then
        if [ $? -eq 1 ]; then
            ((critical_failed++))
        else
            ((warnings++))
        fi
    fi

    if ! check_disk_space; then
        if [ $? -eq 1 ]; then
            ((critical_failed++))
        else
            ((warnings++))
        fi
    fi

    if ! check_system_load; then
        ((warnings++))
    fi
    echo ""

    # WittyPi configuration check
    print_header "Power Management Configuration"

    if ! check_wittypi_power_on_config; then
        if [ $? -eq 1 ]; then
            ((critical_failed++))
        else
            ((warnings++))
        fi
    fi

    # Summary
    echo ""
    print_header "Summary"

    if [ "$critical_failed" -eq 0 ] && [ "$warnings" -eq 0 ]; then
        echo -e "${GREEN}✓ All services healthy!${NC}"
        echo ""
        return 0
    elif [ "$critical_failed" -eq 0 ]; then
        echo -e "${YELLOW}⚠  $warnings warning(s) - review above${NC}"
        echo -e "${GREEN}✓ All critical services operational${NC}"
        echo ""
        return 2
    else
        echo -e "${RED}✗ $critical_failed critical failure(s)${NC}"
        if [ "$warnings" -gt 0 ]; then
            echo -e "${YELLOW}⚠  $warnings warning(s)${NC}"
        fi
        echo ""
        echo -e "${RED}Action required: Fix critical services before deployment${NC}"
        echo ""
        return 1
    fi
}

# ============================================================================
# Entry Point
# ============================================================================

# Check if running as root
if [ "$EUID" -ne 0 ]; then
    echo -e "${YELLOW}Warning: Running without root privileges${NC}"
    echo -e "${YELLOW}Some checks may not work properly. Consider running with sudo.${NC}"
    echo ""
fi

main
exit $?
