#!/bin/bash

# Stable Diffusion REST Server Installation Script
# This script installs the server as a systemd service

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"
DEFAULT_USER="sd-rest"
CREATE_USER=false
ENABLE_FILE_LOGGING=false
LOG_DIR="/var/log/stable-diffusion-rest"
LOG_FILE="${LOG_DIR}/server.log"
QUEUE_DIR="/var/lib/stable-diffusion-rest/queue"
OUTPUT_DIR="/var/lib/stable-diffusion-rest/output"
MODELS_DIR=""
USER_SPECIFIED=""
NO_SYSTEM_USER=false
CURRENT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# 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
            --user)
                USER_SPECIFIED="$2"
                shift 2
                ;;
            --no-system-user)
                NO_SYSTEM_USER=true
                shift
                ;;
            --models-dir)
                MODELS_DIR="$2"
                shift 2
                ;;
            --install-dir)
                INSTALL_DIR="$2"
                shift 2
                ;;
            --help)
                echo "Usage: $0 [options]"
                echo ""
                echo "Options:"
                echo "  --user <username>          Specify system user to run the service (default: sd-rest)"
                echo "  --no-system-user           Do not ask for user, use current user"
                echo "  --models-dir <path>        Path to models directory (required)"
                echo "  --install-dir <path>       Installation directory (default: /opt/stable-diffusion-rest)"
                echo "  --help                     Show this help message"
                echo ""
                echo "Example:"
                echo "  sudo $0 --user sd-rest --models-dir /data/SD_MODELS"
                exit 0
                ;;
            *)
                print_error "Unknown option: $1"
                echo "Use --help for usage information"
                exit 1
                ;;
        esac
    done
}

# Function to ask for system user
ask_for_user() {
    if [[ -n "$USER_SPECIFIED" ]]; then
        DEFAULT_USER="$USER_SPECIFIED"
    elif [[ "$NO_SYSTEM_USER" == true ]]; then
        DEFAULT_USER="$SUDO_USER"
        if [[ -z "$DEFAULT_USER" ]]; then
            DEFAULT_USER=$(whoami)
        fi
        print_info "Using current user: $DEFAULT_USER"
        return
    fi

    echo ""
    read -p "Enter system user to run the service [$DEFAULT_USER]: " input_user
    if [[ -n "$input_user" ]]; then
        DEFAULT_USER="$input_user"
    fi

    # Check if user exists
    if ! id "$DEFAULT_USER" &>/dev/null; then
        print_warning "User '$DEFAULT_USER' does not exist"
        read -p "Create user '$DEFAULT_USER'? (y/n): " create_choice
        if [[ "$create_choice" =~ ^[Yy]$ ]]; then
            CREATE_USER=true
        else
            print_error "Cannot proceed without a valid user"
            exit 1
        fi
    fi
}

# Function to ask for models directory
ask_for_models_dir() {
    if [[ -z "$MODELS_DIR" ]]; then
        echo ""
        read -p "Enter path to models directory: " MODELS_DIR
        if [[ -z "$MODELS_DIR" ]]; then
            print_error "Models directory is required"
            exit 1
        fi
    fi

    if [[ ! -d "$MODELS_DIR" ]]; then
        print_error "Models directory does not exist: $MODELS_DIR"
        exit 1
    fi
}

# Function to ask for file logging
ask_for_file_logging() {
    echo ""
    read -p "Enable file-based logging? (y/n) [n]: " logging_choice
    if [[ "$logging_choice" =~ ^[Yy]$ ]]; then
        ENABLE_FILE_LOGGING=true
        read -p "Enter log file path [$LOG_FILE]: " input_log
        if [[ -n "$input_log" ]]; then
            LOG_FILE="$input_log"
        fi
        LOG_DIR=$(dirname "$LOG_FILE")
    fi
}

# Function to create system user
create_system_user() {
    if [[ "$CREATE_USER" == true ]]; then
        print_info "Creating system user: $DEFAULT_USER"
        useradd --system --no-create-home --shell /sbin/nologin "$DEFAULT_USER" || true
    fi
}

# Function to create directories
create_directories() {
    print_info "Creating directories..."

    mkdir -p "$INSTALL_DIR"
    mkdir -p "$QUEUE_DIR"
    mkdir -p "$OUTPUT_DIR"

    if [[ "$ENABLE_FILE_LOGGING" == true ]]; then
        mkdir -p "$LOG_DIR"
        chown -R "$DEFAULT_USER:$DEFAULT_USER" "$LOG_DIR"
    fi

    chown -R "$DEFAULT_USER:$DEFAULT_USER" "$QUEUE_DIR"
    chown -R "$DEFAULT_USER:$DEFAULT_USER" "$OUTPUT_DIR"
}

# Function to copy binary and files
install_files() {
    print_info "Installing server binary..."

    # Check if binary exists in build directory
    if [[ ! -f "$CURRENT_DIR/build/src/stable-diffusion-rest-server" ]]; then
        print_error "Server binary not found. Please build the project first:"
        print_error "  mkdir -p build && cd build"
        print_error "  cmake .."
        print_error "  cmake --build . --parallel"
        exit 1
    fi

    # Copy binary
    cp "$CURRENT_DIR/build/src/stable-diffusion-rest-server" "$INSTALL_DIR/"
    chmod +x "$INSTALL_DIR/stable-diffusion-rest-server"

    # Copy sd binary if it exists
    if [[ -f "$CURRENT_DIR/build/stable-diffusion.cpp-install/bin/sd" ]]; then
        mkdir -p "$INSTALL_DIR/bin"
        cp "$CURRENT_DIR/build/stable-diffusion.cpp-install/bin/sd" "$INSTALL_DIR/bin/"
        chmod +x "$INSTALL_DIR/bin/sd"
    fi

    chown -R "$DEFAULT_USER:$DEFAULT_USER" "$INSTALL_DIR"
}

# Function to generate systemd service file
generate_service_file() {
    print_info "Generating systemd service file..."

    # Build logging options
    LOGGING_OPTIONS=""
    if [[ "$ENABLE_FILE_LOGGING" == true ]]; then
        LOGGING_OPTIONS="--enable-file-logging --log-file $LOG_FILE"
    fi

    # Replace placeholders in template
    sed -e "s|{{USER}}|$DEFAULT_USER|g" \
        -e "s|{{INSTALL_DIR}}|$INSTALL_DIR|g" \
        -e "s|{{MODELS_DIR}}|$MODELS_DIR|g" \
        -e "s|{{QUEUE_DIR}}|$QUEUE_DIR|g" \
        -e "s|{{OUTPUT_DIR}}|$OUTPUT_DIR|g" \
        -e "s|{{LOG_DIR}}|$LOG_DIR|g" \
        -e "s|{{LOGGING_OPTIONS}}|$LOGGING_OPTIONS|g" \
        "$CURRENT_DIR/stable-diffusion-rest.service.template" > "$SERVICE_FILE"

    print_info "Service file created: $SERVICE_FILE"
}

# Function to enable and start service
setup_service() {
    print_info "Reloading systemd daemon..."
    systemctl daemon-reload

    echo ""
    read -p "Enable service to start on boot? (y/n) [y]: " enable_choice
    if [[ -z "$enable_choice" || "$enable_choice" =~ ^[Yy]$ ]]; then
        systemctl enable "$SERVICE_NAME"
        print_info "Service enabled"
    fi

    echo ""
    read -p "Start service now? (y/n) [y]: " start_choice
    if [[ -z "$start_choice" || "$start_choice" =~ ^[Yy]$ ]]; then
        systemctl start "$SERVICE_NAME"
        print_info "Service started"

        # Show status
        sleep 2
        systemctl status "$SERVICE_NAME" --no-pager -l
    fi
}

# Main installation flow
main() {
    echo "============================================"
    echo "  Stable Diffusion REST Server Installer"
    echo "============================================"
    echo ""

    check_root
    parse_args "$@"
    ask_for_user
    ask_for_models_dir
    ask_for_file_logging

    echo ""
    print_info "Installation Summary:"
    print_info "  User:          $DEFAULT_USER"
    print_info "  Install Dir:   $INSTALL_DIR"
    print_info "  Models Dir:    $MODELS_DIR"
    print_info "  Queue Dir:     $QUEUE_DIR"
    print_info "  Output Dir:    $OUTPUT_DIR"
    print_info "  File Logging:  $ENABLE_FILE_LOGGING"
    if [[ "$ENABLE_FILE_LOGGING" == true ]]; then
        print_info "  Log File:      $LOG_FILE"
    fi
    echo ""

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

    create_system_user
    create_directories
    install_files
    generate_service_file
    setup_service

    echo ""
    print_info "============================================"
    print_info "  Installation Complete!"
    print_info "============================================"
    echo ""
    print_info "Service commands:"
    print_info "  Start:   systemctl start $SERVICE_NAME"
    print_info "  Stop:    systemctl stop $SERVICE_NAME"
    print_info "  Restart: systemctl restart $SERVICE_NAME"
    print_info "  Status:  systemctl status $SERVICE_NAME"
    print_info "  Logs:    journalctl -u $SERVICE_NAME -f"
    echo ""
}

main "$@"