#!/bin/bash

VERSION="1.0.0"
SPINNER="/-\|"

###################
# Color Definitions
###################
NC="\033[00m"
GREEN="\033[01;32m"
RED="\033[01;31m"
YELLOW="\033[01;33m"

###################
# Utility Functions
###################

print_header() {
    echo "
    ____                  ____              
   / __ \_________  _  __/ __ \____  _____ 
  / /_/ / ___/ __ \| |/_/ / / / __ \/ ___/ 
 / ____/ /  / /_/ />  </ /_/ / /_/ / /__   
/_/   /_/   \____/_/|_/_____/\____/\___/   
                                           
 The Proxmox System Doctor
 Version ${VERSION}
======================================"
    echo "Started at: $(date '+%Y-%m-%d %H:%M:%S')"
    echo "Examining patient: $(hostname)"
    echo "======================================"
}


handle_error() {
    echo -e "${RED}Error: $1${NC}"
    exit 1
}

log_message() {
    local level=$1
    local message=$2
    case $level in
        info)    echo -e "${GREEN}[INFO]${NC} $message" ;;
        warn)    echo -e "${YELLOW}[WARN]${NC} $message" ;;
        error)   echo -e "${RED}[ERROR]${NC} $message" ;;
    esac
}

show_progress() {
    local pid=$1
    local delay=0.1
    local spinstr='|/-\'
    while ps -p $pid > /dev/null; do
        local temp=${spinstr#?}
        printf " [%c]  " "$spinstr"
        local spinstr=$temp${spinstr%"$temp"}
        sleep $delay
        printf "\b\b\b\b\b\b"
    done
    printf "    \b\b\b\b"
}

check_requirements() {
    log_message info "Checking medical equipment..."
    local tools=("dmidecode" "lscpu" "ip" "smartctl" "sensors" "netstat")
    for tool in "${tools[@]}"; do
        if ! command -v "$tool" >/dev/null 2>&1; then
            handle_error "Required instrument '$tool' is missing"
        fi
    done
}

checkIfOnHypervisor() {
    if ! command -v pveversion >/dev/null 2>&1; then
        return 1
    fi
    return 0
}

###################
# System Information Functions
###################
get_system_info() {
    echo -e "\n${GREEN}=== System Information ===${NC}"
    echo -e "\n${GREEN}=== Diagnostic Run: $(date '+%Y-%m-%d %H:%M:%S') ===${NC}"
    echo -e "${GREEN}Hostname:$(uname -n)${NC}"
    echo -e "${GREEN}Kernel:$(uname -r)${NC}"
    echo -e "\n${GREEN}=== Proxmox Version ===${NC}"
    pveversion || echo "Not available"
}

get_temp_info() {
    echo -e "\n${GREEN}=== Temperature Information ===${NC}"
    if command -v sensors >/dev/null 2>&1; then
        sensors
    else
        log_message warn "sensors command not found. Install lm-sensors package for temperature monitoring"
    fi
}

get_disk_health() {
    echo -e "\n${GREEN}=== Disk Health Status ===${NC}"
    if command -v smartctl >/dev/null 2>&1; then
        for disk in $(lsblk -d -o name | grep -E '^sd|^nvme'); do
            echo -e "\nChecking /dev/$disk:"
            smartctl -H /dev/$disk
        done
    else
        log_message warn "smartctl not found. Install smartmontools for disk health monitoring"
    fi
}

get_cpu_info() {
    cpu_info=$(grep -m 1 -w 'model name' /proc/cpuinfo | awk -F: '{print $2}' | xargs) || {
        echo -e "${RED}Failed to retrieve CPU model information.${NC}"
    }
    cpu_cores=$(lscpu | grep '^CPU(s):' | awk '{print $2}')
    cpu_mhz=$(lscpu | grep 'MHz' | awk '{print $4}')
    
    echo -e "${GREEN}CPU Model:${NC} $cpu_info"
    echo -e "${GREEN}CPU Cores:${NC} $cpu_cores"
    echo -e "${GREEN}CPU MHz:${NC} $cpu_mhz"
}

get_ram_info() {
    ram_total=$(free -h | grep 'Mem:' | awk '{print $2}')
    ram_used=$(free -h | grep 'Mem:' | awk '{print $3}')
    ram_free=$(free -h | grep 'Mem:' | awk '{print $4}')
    
    echo -e "${GREEN}Total RAM:${NC} $ram_total"
    echo -e "${GREEN}Used RAM:${NC} $ram_used"
    echo -e "${GREEN}Free RAM:${NC} $ram_free"
}

get_storage_info() {
    echo -e "${GREEN}Storage Information:${NC}"
    df -h --output=source,size,used,avail,pcent | grep '^/dev'
    
    if command -v zpool >/dev/null 2>&1; then
        echo -e "\n${GREEN}=== ZFS Pool Status ===${NC}"
        zpool status
    fi
}

get_network_info() {
    default_gateway=$(ip route | grep default | awk '{print $3}')
    ip_addresses=$(hostname -I | xargs)
    echo -e "${GREEN}Default Gateway:${NC} $default_gateway"
    echo -e "${GREEN}IP Addresses:${NC} $ip_addresses"
}

get_detailed_network() {
    echo -e "\n${GREEN}=== Network Interface Details ===${NC}"
    ip -s link show
    echo -e "\n${GREEN}=== Network Statistics ===${NC}"
    netstat -i
}

get_hardware_info() {
    echo -e "${GREEN}BIOS Version:${NC} $(dmidecode -s bios-version)"
    echo -e "\n${GREEN}=== PCI Devices ===${NC}"
    lspci | grep -i -E "vga|ethernet|raid"
}

get_system_status() {
    echo -e "\n${GREEN}=== System Load ===${NC}"
    uptime
    
    echo -e "\n${GREEN}=== Service Status ===${NC}"
    systemctl list-units --type=service --state=running | wc -l
    
    echo -e "\n${GREEN}=== Recent System Errors ===${NC}"
    journalctl -p err -n 5 --no-pager
}

###################
# Proxmox Specific Functions
###################
check_services() {
    echo -e "${GREEN}Checking critical services:${NC}"
    services=("pvedaemon" "pveproxy" "pvecluster" "pve-cluster" "corosync")
    for service in "${services[@]}"; do
        status=$(systemctl is-active "$service")
        echo -e "${GREEN}$service:${NC} $status"
    done
}

check_pve_version() {
    local min_version="6.0"
    local current_version=$(pveversion | grep -oP 'pve-manager/\K[0-9]+\.[0-9]+' || echo "0.0")
    if (( $(echo "$current_version < $min_version" | bc -l) )); then
        log_message warn "Proxmox VE version $current_version may not support all features"
    fi
}

list_vms() {
    if command -v qm >/dev/null 2>&1; then
        echo -e "${GREEN}Virtual Machine Status:${NC}"
        qm list
    else
        handle_error "qm command not found"
    fi
}

list_containers() {
    if command -v pct >/dev/null 2>&1; then
        echo -e "\n${GREEN}=== LXC Container Status ===${NC}"
        pct list
    else
        log_message warn "pct command not found"
    fi
}

###################
# Command Line Interface Functions
###################
help() {
    echo "ProxDoc - The Proxmox System Doctor v${VERSION}"
    echo "Usage: $0 [OPTIONS]"
    echo "A comprehensive diagnostic tool for Proxmox server health checks."
    echo ""
    echo "Treatment Options:"
    echo "  --help       Show this prescription guide"
    echo "  --diags      Perform full system examination"
    echo "  --connect    Make a house call to a remote Proxmox host"
    echo "  --services   Check vital services"
    echo "  --vm-list    Check VM vitals"
    echo "  --backup     Review backup health"
    echo "  --save       Save examination results to medical record"
    exit 0
}


connectToHost() {
    read -rp "Enter the address for house call (hostname/IP): " remote_host
    echo -e "${YELLOW}Preparing for remote consultation with $remote_host...${NC}"
    if ping -c 1 "$remote_host" >/dev/null 2>&1; then
        ssh "root@$remote_host"
    else
        handle_error "Patient $remote_host is not responding"
    fi
}

###################
# Main Functions
###################
runDiags() {
    log_message info "Beginning system examination..."
    (
        get_system_info
        get_cpu_info
        get_ram_info
        get_storage_info
        get_disk_health
        get_network_info
        get_detailed_network
        get_hardware_info
        get_temp_info
        get_system_status
        list_containers
    ) & show_progress $!
    log_message info "Examination complete"
}


checkForInput() {
    case $1 in
        --help)     help ;;
        --diags)    check_requirements; runDiags ;;
        --connect)  connectToHost ;;
        --services) check_services ;;
        --vm-list)  list_vms ;;
        --backup)   echo -e "${GREEN}Backup Status:${NC}"; pvesm status ;;
        --save)     exec 1> >(tee "proxmox_diag_$(date '+%Y%m%d_%H%M%S').log"); runDiags ;;
        *)         echo -e "${RED}Invalid option: $1${NC}"; help ;;
    esac
}

main() {
    print_header

    trap 'echo -e "${RED}Script interrupted.${NC}"; exit 1' INT TERM

    if [[ $EUID -ne 0 ]]; then
        handle_error "This script must be run as root"
    fi

    if checkIfOnHypervisor; then
        runDiags
    else
        connectToHost
    fi
}

###################
# Script Execution
###################
argOne=$1

if [[ -n $argOne ]]; then
    checkForInput "$argOne"
else
    echo "Please enter an option:"
    read -r argOne
    checkForInput "$argOne"
fi

main