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

# ============================================================
# shallot-doctor — Diagnose, backup, and repair Security Shallots
# ============================================================
# Usage:
#   sudo shallot-doctor check          # run all diagnostics
#   sudo shallot-doctor backup          # backup config + DB + certs
#   sudo shallot-doctor restore <file>  # restore from backup tarball
#   sudo shallot-doctor fix-tls         # regenerate self-signed TLS cert
#   sudo shallot-doctor fix-crons       # reinstall maintenance crons
#   sudo shallot-doctor fix-service     # reinstall + restart systemd service
#   sudo shallot-doctor fix-perms       # fix file permissions
#   sudo shallot-doctor reset-config    # reset config.yaml from example
#   sudo shallot-doctor nuke            # full uninstall (prompts for confirmation)

# --------------- Colors & formatting -----------------------
RED='\033[0;31m'
GRN='\033[0;32m'
YLW='\033[1;33m'
BLU='\033[0;34m'
CYN='\033[0;36m'
WHT='\033[1;37m'
DIM='\033[2m'
RST='\033[0m'

ok()   { echo -e "${GRN}  [OK]${RST}  $*"; }
warn() { echo -e "${YLW}  [!!]${RST}  $*"; }
err()  { echo -e "${RED}  [XX]${RST}  $*"; }
info() { echo -e "${BLU}  [>>]${RST}  $*"; }
step() { echo -e "\n${CYN}━━━ $* ${RST}"; }
die()  { err "$*"; exit 1; }

# --------------- Paths ------------------------------------
# Detect install location
if [[ -f /etc/shallots/config.yaml ]]; then
    CONFIG_DIR="/etc/shallots"
elif [[ -f "$HOME/security-shallots/config.yaml" ]]; then
    CONFIG_DIR="$HOME/security-shallots"
elif [[ -f ./config.yaml ]]; then
    CONFIG_DIR="$(pwd)"
else
    CONFIG_DIR=""
fi

INSTALL_DIR="${CONFIG_DIR}"
SERVICE_NAME="shallotd"
BACKUP_DIR="${CONFIG_DIR:+${CONFIG_DIR}/backups}"

# --------------- Banner ------------------------------------
print_banner() {
    echo -e "${CYN}"
    cat <<'BANNER'
   ____  _           _ _       _     ____             _
  / ___|| |__   __ _| | | ___ | |_  |  _ \  ___   ___| |_ ___  _ __
  \___ \| '_ \ / _` | | |/ _ \| __| | | | |/ _ \ / __| __/ _ \| '__|
   ___) | | | | (_| | | | (_) | |_  | |_| | (_) | (__| || (_) | |
  |____/|_| |_|\__,_|_|_|\___/ \__| |____/ \___/ \___|\__\___/|_|
BANNER
    echo -e "${RST}"
}

# --------------- Helper: find config ----------------------
require_config() {
    if [[ -z "$CONFIG_DIR" ]]; then
        die "Cannot find config.yaml. Run from the shallots directory or ensure /etc/shallots/config.yaml exists."
    fi
}

# ============================================================
# CHECK — full diagnostic scan
# ============================================================
cmd_check() {
    local issues=0

    step "Configuration"
    if [[ -n "$CONFIG_DIR" ]]; then
        ok "Config found: ${CONFIG_DIR}/config.yaml"
    else
        err "config.yaml not found"
        issues=$((issues + 1))
    fi

    if [[ -n "$CONFIG_DIR" ]] && grep -q 'tls_cert:' "$CONFIG_DIR/config.yaml" 2>/dev/null; then
        local cert_path
        cert_path=$(grep 'tls_cert:' "$CONFIG_DIR/config.yaml" | awk '{print $2}' | tr -d '"')
        if [[ -n "$cert_path" && -f "$cert_path" ]]; then
            # Check expiry
            local expiry
            expiry=$(openssl x509 -enddate -noout -in "$cert_path" 2>/dev/null | cut -d= -f2)
            if [[ -n "$expiry" ]]; then
                if openssl x509 -checkend 604800 -noout -in "$cert_path" 2>/dev/null; then
                    ok "TLS cert valid (expires: $expiry)"
                else
                    warn "TLS cert expires within 7 days: $expiry"
                    issues=$((issues + 1))
                fi
            fi
        elif [[ -n "$cert_path" ]]; then
            err "TLS cert not found: $cert_path"
            issues=$((issues + 1))
        fi
    fi

    step "Services"
    # shallotd
    if systemctl is-active "$SERVICE_NAME" &>/dev/null; then
        ok "$SERVICE_NAME: active"
    elif systemctl is-enabled "$SERVICE_NAME" &>/dev/null; then
        err "$SERVICE_NAME: enabled but NOT running"
        issues=$((issues + 1))
    else
        err "$SERVICE_NAME: not installed as service"
        issues=$((issues + 1))
    fi

    # Suricata
    if systemctl is-active suricata &>/dev/null; then
        ok "suricata: active"
    elif command -v suricata &>/dev/null; then
        warn "suricata: installed but not running"
        issues=$((issues + 1))
    else
        warn "suricata: not installed"
    fi

    # Wazuh Manager
    if systemctl is-active wazuh-manager &>/dev/null; then
        ok "wazuh-manager: active"
        # Count agents
        local agent_count
        agent_count=$(/var/ossec/bin/manage_agents -l 2>/dev/null | grep -c "ID:" || echo 0)
        info "  Enrolled agents: $agent_count"
    elif dpkg -s wazuh-manager &>/dev/null 2>&1 || rpm -q wazuh-manager &>/dev/null 2>&1; then
        err "wazuh-manager: installed but not running"
        issues=$((issues + 1))
    else
        warn "wazuh-manager: not installed (server only)"
    fi

    # Wazuh Agent (endpoint only)
    if systemctl is-active wazuh-agent &>/dev/null; then
        ok "wazuh-agent (Clove): active"
        if [[ -s /var/ossec/etc/client.keys ]]; then
            local agent_id
            agent_id=$(awk '{print $1}' /var/ossec/etc/client.keys | head -1)
            info "  Agent ID: $agent_id"
        fi
        if grep -q "Connected to the server" /var/ossec/logs/ossec.log 2>/dev/null; then
            ok "  Connected to manager"
        else
            warn "  Not connected to manager (check ossec.log)"
            issues=$((issues + 1))
        fi
    fi

    # CrowdSec
    if systemctl is-active crowdsec &>/dev/null; then
        ok "crowdsec: active"
    fi

    step "Ports"
    for port_desc in "8844:dashboard" "1514:wazuh-events" "1515:wazuh-enrollment" "8855:argus-webhook"; do
        local port="${port_desc%%:*}"
        local desc="${port_desc##*:}"
        if ss -tlnp 2>/dev/null | grep -q ":${port} "; then
            ok "Port $port ($desc): listening"
        else
            # Only warn for dashboard, others may be intentionally off
            if [[ "$port" == "8844" ]]; then
                err "Port $port ($desc): NOT listening"
                issues=$((issues + 1))
            fi
        fi
    done

    step "Database"
    if [[ -n "$CONFIG_DIR" ]]; then
        local db_path
        db_path=$(grep -A1 'storage:' "$CONFIG_DIR/config.yaml" 2>/dev/null | grep 'db_path:' | awk '{print $2}' | tr -d '"')
        # Fallback to common locations
        [[ -z "$db_path" ]] && db_path="${CONFIG_DIR}/shallots.db"
        if [[ -f "$db_path" ]]; then
            local db_size
            db_size=$(du -h "$db_path" | awk '{print $1}')
            ok "Database: $db_path ($db_size)"
            # Quick integrity check
            if command -v sqlite3 &>/dev/null; then
                local integrity
                integrity=$(sqlite3 "$db_path" "PRAGMA integrity_check" 2>/dev/null | head -1)
                if [[ "$integrity" == "ok" ]]; then
                    ok "  Integrity check: passed"
                else
                    err "  Integrity check: FAILED"
                    issues=$((issues + 1))
                fi
                local alert_count
                alert_count=$(sqlite3 "$db_path" "SELECT COUNT(*) FROM alerts" 2>/dev/null || echo "?")
                info "  Alerts in DB: $alert_count"
            fi
        else
            warn "Database not found at expected path: $db_path"
        fi
    fi

    # Check backups (look for both tarball and daily .db backups)
    if [[ -n "$BACKUP_DIR" && -d "$BACKUP_DIR" ]]; then
        local latest_backup
        latest_backup=$(find "$BACKUP_DIR" -maxdepth 1 \( -name "shallots-backup-*.tar.gz" -o -name "shallots-*.db" \) -printf '%T@ %p\n' 2>/dev/null | sort -rn | head -1 | awk '{print $2}')
        if [[ -n "$latest_backup" && -f "$latest_backup" ]]; then
            local backup_age
            backup_age=$(( ( $(date +%s) - $(stat -c %Y "$latest_backup" 2>/dev/null || echo 0) ) / 86400 ))
            ok "Latest backup: $(basename "$latest_backup") (${backup_age}d ago)"
            if [[ $backup_age -gt 7 ]]; then
                warn "  Backup is over 7 days old — run: shallot-doctor backup"
                issues=$((issues + 1))
            fi
        else
            warn "No backups found — run: shallot-doctor backup"
            issues=$((issues + 1))
        fi
    fi

    step "Log files"
    for logfile in "/var/log/suricata/eve.json" "/var/ossec/logs/alerts/alerts.json"; do
        if [[ -f "$logfile" ]]; then
            if [[ -r "$logfile" ]]; then
                ok "$logfile (readable)"
            else
                err "$logfile (exists but NOT readable — fix permissions)"
                issues=$((issues + 1))
            fi
        fi
    done

    step "Cron jobs"
    local has_cron=false
    # Check both root and regular user crontabs
    local all_crons=""
    all_crons+=$(crontab -l 2>/dev/null || true)
    all_crons+=$(crontab -u root -l 2>/dev/null || true)
    # Also check common service user
    for u in om shallots; do
        all_crons+=$(crontab -u "$u" -l 2>/dev/null || true)
    done
    if echo "$all_crons" | grep -qE "suricata-update|backup\.sh|shallot"; then
        ok "Maintenance crons found"
        has_cron=true
    fi
    if [[ "$has_cron" == false ]]; then
        warn "No maintenance crons detected — run: shallot-doctor fix-crons"
        issues=$((issues + 1))
    fi

    # Summary
    echo ""
    if [[ $issues -eq 0 ]]; then
        echo -e "${GRN}  All checks passed. System is healthy.${RST}"
    else
        echo -e "${YLW}  Found $issues issue(s). Run specific fix commands to resolve:${RST}"
        echo -e "${DIM}    shallot-doctor fix-tls       # regenerate TLS certificate${RST}"
        echo -e "${DIM}    shallot-doctor fix-crons     # reinstall maintenance crons${RST}"
        echo -e "${DIM}    shallot-doctor fix-service   # reinstall systemd service${RST}"
        echo -e "${DIM}    shallot-doctor fix-perms     # fix file permissions${RST}"
        echo -e "${DIM}    shallot-doctor backup        # create a backup now${RST}"
    fi
    echo ""
}

# ============================================================
# BACKUP — snapshot config + DB + certs
# ============================================================
cmd_backup() {
    require_config

    mkdir -p "$BACKUP_DIR"
    local timestamp
    timestamp=$(date +%Y%m%d-%H%M%S)
    local tarball="${BACKUP_DIR}/shallots-backup-${timestamp}.tar.gz"

    local files_to_backup=()

    # Config
    [[ -f "$CONFIG_DIR/config.yaml" ]] && files_to_backup+=("$CONFIG_DIR/config.yaml")

    # TLS certs
    if grep -q 'tls_cert:' "$CONFIG_DIR/config.yaml" 2>/dev/null; then
        local cert_path key_path
        cert_path=$(grep 'tls_cert:' "$CONFIG_DIR/config.yaml" | awk '{print $2}' | tr -d '"')
        key_path=$(grep 'tls_key:' "$CONFIG_DIR/config.yaml" | awk '{print $2}' | tr -d '"')
        [[ -f "$cert_path" ]] && files_to_backup+=("$cert_path")
        [[ -f "$key_path" ]] && files_to_backup+=("$key_path")
    fi

    # Database
    local db_path
    db_path=$(grep -A1 'storage:' "$CONFIG_DIR/config.yaml" 2>/dev/null | grep 'db_path:' | awk '{print $2}' | tr -d '"')
    [[ -z "$db_path" ]] && db_path="${CONFIG_DIR}/shallots.db"
    if [[ -f "$db_path" ]] && command -v sqlite3 &>/dev/null; then
        local db_backup="${BACKUP_DIR}/shallots-${timestamp}.db"
        info "Creating safe database copy..."
        sqlite3 "$db_path" ".backup '${db_backup}'"
        files_to_backup+=("$db_backup")
    elif [[ -f "$db_path" ]]; then
        files_to_backup+=("$db_path")
    fi

    # Backup script and crontab
    [[ -f "$CONFIG_DIR/backup.sh" ]] && files_to_backup+=("$CONFIG_DIR/backup.sh")

    # Suricata threshold
    [[ -f /etc/suricata/threshold.config ]] && files_to_backup+=("/etc/suricata/threshold.config")

    # Custom Suricata rules
    [[ -f /var/lib/suricata/rules/shallots-home.rules ]] && files_to_backup+=("/var/lib/suricata/rules/shallots-home.rules")

    # Systemd service
    [[ -f /etc/systemd/system/shallotd.service ]] && files_to_backup+=("/etc/systemd/system/shallotd.service")

    # Crontabs
    crontab -l > "${BACKUP_DIR}/crontab-${USER}-${timestamp}.txt" 2>/dev/null || true
    files_to_backup+=("${BACKUP_DIR}/crontab-${USER}-${timestamp}.txt")
    if [[ $EUID -eq 0 ]]; then
        crontab -u root -l > "${BACKUP_DIR}/crontab-root-${timestamp}.txt" 2>/dev/null || true
        files_to_backup+=("${BACKUP_DIR}/crontab-root-${timestamp}.txt")
    fi

    # Create tarball
    info "Packing ${#files_to_backup[@]} files..."
    tar czf "$tarball" "${files_to_backup[@]}" 2>/dev/null || \
    tar czf "$tarball" --ignore-failed-read "${files_to_backup[@]}" 2>/dev/null

    # Cleanup temp files
    rm -f "${BACKUP_DIR}/shallots-${timestamp}.db" \
          "${BACKUP_DIR}/crontab-"*"-${timestamp}.txt" 2>/dev/null

    local size
    size=$(du -h "$tarball" | awk '{print $1}')
    ok "Backup created: $tarball ($size)"

    # Prune old backups (keep last 10)
    local count
    count=$(ls -1 "$BACKUP_DIR"/shallots-backup-*.tar.gz 2>/dev/null | wc -l)
    if [[ $count -gt 10 ]]; then
        ls -t "$BACKUP_DIR"/shallots-backup-*.tar.gz | tail -n +11 | xargs rm -f
        info "Pruned old backups (kept 10)"
    fi
}

# ============================================================
# RESTORE — restore from backup tarball
# ============================================================
cmd_restore() {
    local tarball="${1:-}"
    if [[ -z "$tarball" || ! -f "$tarball" ]]; then
        die "Usage: shallot-doctor restore <backup-file.tar.gz>"
    fi

    step "Restoring from $(basename "$tarball")"

    info "Contents:"
    tar tzf "$tarball" | head -20

    echo ""
    echo -en "  ${YLW}This will overwrite current config and database. Continue? [y/N]: ${RST}"
    read -r confirm
    if [[ "$confirm" != "y" && "$confirm" != "Y" ]]; then
        info "Cancelled."
        return
    fi

    # Stop service first
    if systemctl is-active "$SERVICE_NAME" &>/dev/null; then
        info "Stopping $SERVICE_NAME..."
        systemctl stop "$SERVICE_NAME"
    fi

    # Extract (preserving paths)
    tar xzf "$tarball" -C / 2>/dev/null || tar xzf "$tarball" 2>/dev/null
    ok "Files restored"

    # Restart
    if systemctl is-enabled "$SERVICE_NAME" &>/dev/null; then
        systemctl start "$SERVICE_NAME"
        sleep 3
        if systemctl is-active "$SERVICE_NAME" &>/dev/null; then
            ok "$SERVICE_NAME restarted"
        else
            err "$SERVICE_NAME failed to start — check: journalctl -u $SERVICE_NAME -n 30"
        fi
    fi
}

# ============================================================
# FIX-TLS — regenerate self-signed certificate
# ============================================================
cmd_fix_tls() {
    require_config

    local cert_path="${CONFIG_DIR}/tls.cert"
    local key_path="${CONFIG_DIR}/tls.key"

    # Check if custom paths in config
    if grep -q 'tls_cert:' "$CONFIG_DIR/config.yaml" 2>/dev/null; then
        cert_path=$(grep 'tls_cert:' "$CONFIG_DIR/config.yaml" | awk '{print $2}' | tr -d '"')
        key_path=$(grep 'tls_key:' "$CONFIG_DIR/config.yaml" | awk '{print $2}' | tr -d '"')
    fi

    step "Generating self-signed TLS certificate"

    local hostname_val
    hostname_val=$(hostname)

    openssl req -x509 -newkey rsa:2048 -nodes \
        -keyout "$key_path" \
        -out "$cert_path" \
        -days 365 \
        -subj "/CN=${hostname_val}" \
        -addext "subjectAltName=DNS:${hostname_val},DNS:localhost,IP:127.0.0.1" \
        2>/dev/null

    chmod 600 "$key_path"
    chmod 644 "$cert_path"

    ok "Certificate: $cert_path"
    ok "Key: $key_path"
    ok "Valid for 365 days"

    # Add to config if not already there
    if ! grep -q 'tls_cert:' "$CONFIG_DIR/config.yaml" 2>/dev/null; then
        info "Add these to your config.yaml under the 'web:' section:"
        echo -e "  ${DIM}  tls_cert: \"$cert_path\"${RST}"
        echo -e "  ${DIM}  tls_key: \"$key_path\"${RST}"
    fi

    info "Restart shallotd to pick up new cert: systemctl restart $SERVICE_NAME"
}

# ============================================================
# FIX-CRONS — reinstall maintenance crons
# ============================================================
cmd_fix_crons() {
    require_config

    step "Installing maintenance crons"

    # Root crons (suricata update + logrotate)
    if command -v suricata-update &>/dev/null; then
        local root_cron
        root_cron=$(crontab -u root -l 2>/dev/null || echo "")
        if ! echo "$root_cron" | grep -q "suricata-update"; then
            (echo "$root_cron"; echo "0 3 * * * /usr/bin/suricata-update && /usr/bin/suricatasc -c reload-rules >/dev/null 2>&1") | crontab -u root -
            ok "Added: Suricata rule update daily at 3 AM (root crontab)"
        else
            ok "Suricata cron already exists"
        fi
    fi

    # Logrotate for Suricata
    if [[ -d /var/log/suricata ]] && [[ ! -f /etc/logrotate.d/suricata ]]; then
        cat > /etc/logrotate.d/suricata <<'LOGROTATE'
/var/log/suricata/*.json {
    daily
    rotate 7
    compress
    delaycompress
    missingok
    notifempty
    postrotate
        systemctl reload suricata 2>/dev/null || true
    endscript
}
LOGROTATE
        ok "Added: Suricata logrotate (7-day retention)"
    fi

    # DB backup cron
    local backup_script="${CONFIG_DIR}/backup.sh"
    if [[ ! -f "$backup_script" ]]; then
        local db_path
        db_path=$(grep -A1 'storage:' "$CONFIG_DIR/config.yaml" 2>/dev/null | grep 'db_path:' | awk '{print $2}' | tr -d '"')
        [[ -z "$db_path" ]] && db_path="${CONFIG_DIR}/shallots.db"

        cat > "$backup_script" <<BACKUP
#!/bin/bash
BACKUP_DIR="${CONFIG_DIR}/backups"
mkdir -p "\$BACKUP_DIR"
sqlite3 "$db_path" ".backup '\$BACKUP_DIR/shallots-\$(date +%Y%m%d).db'"
find "\$BACKUP_DIR" -name "shallots-*.db" -mtime +7 -delete
BACKUP
        chmod +x "$backup_script"
        ok "Created: $backup_script"
    fi

    # User cron for backup
    local user_cron
    user_cron=$(crontab -l 2>/dev/null || echo "")
    if ! echo "$user_cron" | grep -q "backup.sh"; then
        (echo "$user_cron"; echo "0 4 * * * ${backup_script} >/dev/null 2>&1") | crontab -
        ok "Added: DB backup daily at 4 AM (user crontab)"
    else
        ok "Backup cron already exists"
    fi
}

# ============================================================
# FIX-SERVICE — reinstall systemd service
# ============================================================
cmd_fix_service() {
    require_config

    step "Installing systemd service"

    local service_file="/etc/systemd/system/${SERVICE_NAME}.service"
    local install_dir="$CONFIG_DIR"

    # Find the setup/systemd template
    local template=""
    for candidate in "${install_dir}/setup/systemd/shallot.service" \
                     "$(dirname "$install_dir")/setup/systemd/shallot.service"; do
        if [[ -f "$candidate" ]]; then
            template="$candidate"
            break
        fi
    done

    if [[ -n "$template" ]]; then
        cp "$template" "$service_file"
        ok "Copied service from template: $template"
    else
        # Generate a basic service file
        cat > "$service_file" <<SERVICE
[Unit]
Description=Security Shallots daemon
After=network.target

[Service]
Type=simple
User=$(whoami)
WorkingDirectory=${install_dir}
ExecStart=$(command -v python3 || echo /usr/bin/python3) -m shallots.cli start --config ${install_dir}/config.yaml
Restart=on-failure
RestartSec=10
NoNewPrivileges=yes
PrivateTmp=yes

[Install]
WantedBy=multi-user.target
SERVICE
        ok "Generated service file: $service_file"
    fi

    systemctl daemon-reload
    systemctl enable "$SERVICE_NAME"
    systemctl restart "$SERVICE_NAME"

    sleep 3
    if systemctl is-active "$SERVICE_NAME" &>/dev/null; then
        ok "$SERVICE_NAME is running"
    else
        err "$SERVICE_NAME failed to start"
        warn "Check: journalctl -u $SERVICE_NAME -n 30"
    fi
}

# ============================================================
# FIX-PERMS — fix file permissions
# ============================================================
cmd_fix_perms() {
    require_config

    step "Fixing permissions"

    local current_user
    current_user=$(stat -c '%U' "$CONFIG_DIR/config.yaml" 2>/dev/null || whoami)

    # Config file (readable by owner only — has secrets)
    if [[ -f "$CONFIG_DIR/config.yaml" ]]; then
        chmod 600 "$CONFIG_DIR/config.yaml"
        ok "config.yaml: 600 (owner-only)"
    fi

    # TLS key
    if grep -q 'tls_key:' "$CONFIG_DIR/config.yaml" 2>/dev/null; then
        local key_path
        key_path=$(grep 'tls_key:' "$CONFIG_DIR/config.yaml" | awk '{print $2}' | tr -d '"')
        if [[ -f "$key_path" ]]; then
            chmod 600 "$key_path"
            ok "TLS key: 600"
        fi
    fi

    # Suricata EVE log — add user to suricata group
    if getent group suricata &>/dev/null; then
        if ! id -nG "$current_user" 2>/dev/null | grep -qw suricata; then
            usermod -aG suricata "$current_user"
            ok "Added $current_user to suricata group"
        else
            ok "$current_user already in suricata group"
        fi
    fi

    # Wazuh alerts — add user to wazuh group
    if getent group wazuh &>/dev/null; then
        if ! id -nG "$current_user" 2>/dev/null | grep -qw wazuh; then
            usermod -aG wazuh "$current_user"
            ok "Added $current_user to wazuh group"
        else
            ok "$current_user already in wazuh group"
        fi
    fi

    # Backup directory
    if [[ -d "$BACKUP_DIR" ]]; then
        chmod 700 "$BACKUP_DIR"
        ok "Backup dir: 700"
    fi
}

# ============================================================
# RESET-CONFIG — reset config from example
# ============================================================
cmd_reset_config() {
    require_config

    step "Resetting configuration"

    local example=""
    for candidate in "${CONFIG_DIR}/config.example.yaml" \
                     "$(dirname "$CONFIG_DIR")/config.example.yaml"; do
        if [[ -f "$candidate" ]]; then
            example="$candidate"
            break
        fi
    done

    if [[ -z "$example" ]]; then
        die "config.example.yaml not found. Clone the repo first: git clone https://github.com/benolenick/security-shallots.git"
    fi

    # Backup current config first
    if [[ -f "$CONFIG_DIR/config.yaml" ]]; then
        local backup="${CONFIG_DIR}/config.yaml.bak.$(date +%Y%m%d-%H%M%S)"
        cp "$CONFIG_DIR/config.yaml" "$backup"
        ok "Current config backed up to: $backup"
    fi

    cp "$example" "$CONFIG_DIR/config.yaml"
    ok "Config reset from: $example"
    warn "You will need to re-enter secrets (web password, SMTP, Twilio, AI keys)"
    info "Edit: $CONFIG_DIR/config.yaml"
    info "Then restart: systemctl restart $SERVICE_NAME"
}

# ============================================================
# NUKE — full uninstall
# ============================================================
cmd_nuke() {
    echo ""
    echo -e "${RED}  WARNING: This will remove Security Shallots completely.${RST}"
    echo -e "${RED}  Database, config, crons, systemd service — all gone.${RST}"
    echo ""
    echo -en "  ${YLW}Type 'yes-nuke-it' to confirm: ${RST}"
    read -r confirm
    if [[ "$confirm" != "yes-nuke-it" ]]; then
        info "Cancelled."
        return
    fi

    step "Removing Security Shallots"

    # Stop services
    systemctl stop "$SERVICE_NAME" 2>/dev/null || true
    systemctl disable "$SERVICE_NAME" 2>/dev/null || true
    rm -f "/etc/systemd/system/${SERVICE_NAME}.service"
    systemctl daemon-reload

    # Remove crons
    crontab -l 2>/dev/null | grep -v "backup.sh\|shallot" | crontab - 2>/dev/null || true
    ok "Crons removed"

    ok "Service removed"
    warn "Config, database, and code remain at: ${CONFIG_DIR}"
    info "To fully delete: rm -rf ${CONFIG_DIR}"
}

# ============================================================
# fix-rules — Install safe Suricata rule update wrapper
# ============================================================
cmd_fix_rules() {
    step "Installing safe rule update wrapper"

    cat > /usr/local/bin/shallots-rule-update <<'RULEUPDATE'
#!/usr/bin/env bash
set -euo pipefail
# shallots-rule-update — safe suricata-update with auto-rollback

LOG="/var/log/suricata/rule-update.log"
RULES_DIR="/var/lib/suricata/rules"
BACKUP_DIR="/var/lib/suricata/rules-backup"
CUSTOM_RULES="${RULES_DIR}/shallots-home.rules"

log() { echo "$(date '+%Y-%m-%d %H:%M:%S') $*" | tee -a "$LOG"; }

log "=== Rule update started ==="

# Pre-update: count current rules
pre_count=0
if [ -d "$RULES_DIR" ]; then
    pre_count=$(find "$RULES_DIR" -name "*.rules" -exec grep -c '^alert\|^drop\|^reject\|^pass' {} + 2>/dev/null | tail -1 || echo "0")
fi
log "Pre-update rule count: $pre_count"

# Backup current rules
rm -rf "$BACKUP_DIR"
cp -a "$RULES_DIR" "$BACKUP_DIR"
log "Rules backed up to $BACKUP_DIR"

# Run suricata-update
log "Running suricata-update..."
if ! suricata-update 2>&1 | tee -a "$LOG"; then
    log "ERROR: suricata-update failed, rolling back"
    rm -rf "$RULES_DIR"
    cp -a "$BACKUP_DIR" "$RULES_DIR"
    log "Rollback complete"
    exit 1
fi

# Restore custom rules if they got clobbered
if [ -f "${BACKUP_DIR}/shallots-home.rules" ] && [ ! -f "$CUSTOM_RULES" ]; then
    cp "${BACKUP_DIR}/shallots-home.rules" "$CUSTOM_RULES"
    log "Restored custom shallots-home.rules"
fi

# Post-update: count new rules
post_count=$(find "$RULES_DIR" -name "*.rules" -exec grep -c '^alert\|^drop\|^reject\|^pass' {} + 2>/dev/null | tail -1 || echo "0")
log "Post-update rule count: $post_count"

# Check for >10% rule drop
if [ "$pre_count" -gt 0 ]; then
    threshold=$(( pre_count * 90 / 100 ))
    if [ "$post_count" -lt "$threshold" ]; then
        log "WARNING: Rule count dropped >10% ($pre_count -> $post_count), rolling back"
        rm -rf "$RULES_DIR"
        cp -a "$BACKUP_DIR" "$RULES_DIR"
        log "Rollback complete (rules restored to $pre_count)"
        exit 1
    fi
fi

# Syntax validation
log "Running suricata -T (syntax check)..."
if ! suricata -T -c /etc/suricata/suricata.yaml 2>&1 | tee -a "$LOG"; then
    log "ERROR: Syntax check failed, rolling back"
    rm -rf "$RULES_DIR"
    cp -a "$BACKUP_DIR" "$RULES_DIR"
    log "Rollback complete"
    exit 1
fi

# All good — reload rules
log "Syntax check passed. Reloading rules..."
suricatasc -c reload-rules 2>&1 | tee -a "$LOG" || true
log "=== Rule update completed successfully ($post_count rules loaded) ==="

# Cleanup old backup
rm -rf "$BACKUP_DIR"
RULEUPDATE

    chmod +x /usr/local/bin/shallots-rule-update
    ok "Installed /usr/local/bin/shallots-rule-update"

    # Replace raw suricata-update in cron with the safe wrapper
    if crontab -l 2>/dev/null | grep -q "suricata-update"; then
        info "Replacing raw suricata-update cron with safe wrapper..."
        crontab -l 2>/dev/null | sed 's|suricata-update.*|/usr/local/bin/shallots-rule-update|g' | crontab -
        ok "Cron updated to use safe wrapper"
    else
        # Install cron for daily 3 AM
        local cron_line="0 3 * * * /usr/local/bin/shallots-rule-update"
        (crontab -l 2>/dev/null; echo "$cron_line") | crontab -
        ok "Cron installed: daily at 3 AM"
    fi

    info "Safe rule update features:"
    info "  - Pre-update rule backup"
    info "  - Auto-rollback if rule count drops >10%"
    info "  - Auto-rollback if suricata -T syntax check fails"
    info "  - Custom shallots-home.rules preserved"
    info "  - Logs to /var/log/suricata/rule-update.log"
    echo ""
    info "Run manually: sudo shallots-rule-update"
    info "View log: tail -f /var/log/suricata/rule-update.log"
}

# ============================================================
# Usage
# ============================================================
usage() {
    cat <<EOF
Usage: $(basename "$0") <command>

Commands:
  check           Run full diagnostic scan
  backup          Backup config, DB, certs, and crons
  restore <file>  Restore from a backup tarball
  fix-tls         Regenerate self-signed TLS certificate
  fix-crons       Reinstall maintenance crons (suricata update, DB backup)
  fix-service     Reinstall + restart systemd service
  fix-perms       Fix file permissions (config, logs, groups)
  reset-config    Reset config.yaml from example template
  fix-rules       Install safe rule update wrapper (auto-rollback on failure)
  nuke            Full uninstall (requires confirmation)

Quick recovery after a broken config:
  sudo shallot-doctor reset-config    # start fresh
  sudo shallot-doctor fix-tls         # if HTTPS broke
  sudo shallot-doctor fix-service     # if service won't start

Full rebuild from backup:
  sudo shallot-doctor restore backups/shallots-backup-20260302.tar.gz

Endpoint (Clove) recovery:
  curl -fsSL https://raw.githubusercontent.com/benolenick/security-shallots/main/setup/endpoint/clove \\
    | sudo bash -s -- --manager YOUR_SERVER_IP
EOF
    exit 0
}

# ============================================================
# Main
# ============================================================
main() {
    local cmd="${1:-}"
    shift 2>/dev/null || true

    case "$cmd" in
        check)        print_banner; cmd_check ;;
        backup)       print_banner; cmd_backup ;;
        restore)      print_banner; cmd_restore "$@" ;;
        fix-tls)      print_banner; cmd_fix_tls ;;
        fix-crons)    print_banner; cmd_fix_crons ;;
        fix-service)  print_banner; cmd_fix_service ;;
        fix-perms)    print_banner; cmd_fix_perms ;;
        reset-config) print_banner; cmd_reset_config ;;
        fix-rules)    print_banner; cmd_fix_rules ;;
        nuke)         print_banner; cmd_nuke ;;
        -h|--help|help|"") usage ;;
        *)            die "Unknown command: $cmd (use --help)" ;;
    esac
}

main "$@"
