#!/bin/bash

# Stable Diffusion REST Server Uninstallation Script
# This script removes the systemd service and optionally cleans up files

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Default values
SERVICE_NAME="stable-diffusion-rest"
SERVICE_FILE="/etc/systemd/system/${SERVICE_NAME}.service"
INSTALL_DIR="/opt/stable-diffusion-rest"
LOG_DIR="/var/log/stable-diffusion-rest"
QUEUE_DIR="/var/lib/stable-diffusion-rest/queue"
OUTPUT_DIR="/var/lib/stable-diffusion-rest/output"
DATA_DIR="/var/lib/stable-diffusion-rest"
REMOVE_DATA=false
REMOVE_USER=false
USER_TO_REMOVE=""

# Function to print colored output
print_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Function to check if script is run as root
check_root() {
    if [[ $EUID -ne 0 ]]; then
        print_error "This script must be run as root (use sudo)"
        exit 1
    fi
}

# Function to parse command line arguments
parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --remove-data)
                REMOVE_DATA=true
                shift
                ;;
            --remove-user)
                REMOVE_USER=true
                shift
                ;;
            --help)
                echo "Usage: $0 [options]"
                echo ""
                echo "Options:"
                echo "  --remove-data       Also remove queue, output, and log directories"
                echo "  --remove-user       Also remove the system user (asks for confirmation)"
                echo "  --help              Show this help message"
                echo ""
                echo "Example:"
                echo "  sudo $0"
                echo "  sudo $0 --remove-data --remove-user"
                exit 0
                ;;
            *)
                print_error "Unknown option: $1"
                echo "Use --help for usage information"
                exit 1
                ;;
        esac
    done
}

# Function to stop and disable service
stop_service() {
    if systemctl is-active --quiet "$SERVICE_NAME"; then
        print_info "Stopping service..."
        systemctl stop "$SERVICE_NAME"
    fi

    if systemctl is-enabled --quiet "$SERVICE_NAME" 2>/dev/null; then
        print_info "Disabling service..."
        systemctl disable "$SERVICE_NAME"
    fi
}

# Function to remove service file
remove_service() {
    if [[ -f "$SERVICE_FILE" ]]; then
        print_info "Removing service file..."

        # Extract user from service file before removing
        if [[ -f "$SERVICE_FILE" ]]; then
            USER_TO_REMOVE=$(grep "^User=" "$SERVICE_FILE" | cut -d= -f2)
        fi

        rm -f "$SERVICE_FILE"
        systemctl daemon-reload
    else
        print_warning "Service file not found: $SERVICE_FILE"
    fi
}

# Function to remove installation directory
remove_install_dir() {
    if [[ -d "$INSTALL_DIR" ]]; then
        print_info "Removing installation directory: $INSTALL_DIR"
        rm -rf "$INSTALL_DIR"
    fi
}

# Function to remove data directories
remove_data_dirs() {
    if [[ "$REMOVE_DATA" == true ]]; then
        echo ""
        print_warning "This will remove all queue data, generated outputs, and logs"
        read -p "Are you sure? (y/n): " confirm
        if [[ "$confirm" =~ ^[Yy]$ ]]; then
            if [[ -d "$LOG_DIR" ]]; then
                print_info "Removing log directory: $LOG_DIR"
                rm -rf "$LOG_DIR"
            fi
            if [[ -d "$DATA_DIR" ]]; then
                print_info "Removing data directory: $DATA_DIR"
                rm -rf "$DATA_DIR"
            fi
        fi
    else
        echo ""
        print_info "Data directories preserved:"
        print_info "  Queue:  $QUEUE_DIR"
        print_info "  Output: $OUTPUT_DIR"
        print_info "  Logs:   $LOG_DIR"
        print_info ""
        print_info "To remove these, run: sudo $0 --remove-data"
    fi
}

# Function to remove system user
remove_system_user() {
    if [[ "$REMOVE_USER" == true && -n "$USER_TO_REMOVE" ]]; then
        if id "$USER_TO_REMOVE" &>/dev/null; then
            echo ""
            print_warning "This will remove the system user: $USER_TO_REMOVE"
            read -p "Are you sure? (y/n): " confirm
            if [[ "$confirm" =~ ^[Yy]$ ]]; then
                print_info "Removing system user: $USER_TO_REMOVE"
                userdel "$USER_TO_REMOVE" || print_warning "Failed to remove user (may not be a system account)"
            fi
        fi
    elif [[ "$REMOVE_USER" == true ]]; then
        print_warning "Could not determine which user to remove"
    fi
}

# Main uninstallation flow
main() {
    echo "============================================"
    echo " Stable Diffusion REST Server Uninstaller"
    echo "============================================"
    echo ""

    check_root
    parse_args "$@"

    echo ""
    print_info "Uninstallation Summary:"
    print_info "  Service file:    $SERVICE_FILE"
    print_info "  Install dir:     $INSTALL_DIR"
    if [[ "$REMOVE_DATA" == true ]]; then
        print_info "  Remove data:     YES"
    else
        print_info "  Remove data:     NO (use --remove-data to remove)"
    fi
    if [[ "$REMOVE_USER" == true ]]; then
        print_info "  Remove user:     YES (will ask for confirmation)"
    else
        print_info "  Remove user:     NO (use --remove-user to remove)"
    fi
    echo ""

    read -p "Proceed with uninstallation? (y/n): " proceed
    if [[ ! "$proceed" =~ ^[Yy]$ ]]; then
        print_warning "Uninstallation cancelled"
        exit 0
    fi

    stop_service
    remove_service
    remove_install_dir
    remove_data_dirs
    remove_system_user

    echo ""
    print_info "============================================"
    print_info "  Uninstallation Complete!"
    print_info "============================================"
    echo ""
}

main "$@"