#!/usr/bin/env bash

# ComfyUI Tmux Launcher
#
# This script starts ComfyUI in a detached tmux session. It allows for easy
# configuration of session name, port, listen address, and low VRAM mode.
#
# Usage:
#   ./script.sh [SESSION_NAME] [PORT] [LISTEN_ADDRESS] [LOWVRAM]
#
# Arguments:
#   SESSION_NAME   : Name of the tmux session (default: imagegen)
#   PORT           : Port number for ComfyUI to listen on (default: 8080)
#   LISTEN_ADDRESS : IP address to bind to (default: 127.0.0.1)
#   LOWVRAM        : Enable low VRAM mode (default: false)
#
# Examples:
#   1. Run with default settings:
#      ./script.sh
#
#   2. Custom session name:
#      ./script.sh mycustomsession
#
#   3. Custom session, port, and address:
#      ./script.sh mycustomsession 8090 0.0.0.0
#
#   4. Enable low VRAM mode:
#      ./script.sh imagegen 8080 127.0.0.1 true
#
# Note: This script requires tmux to be installed and a virtual environment
# to be set up in the .venv directory.

# Exit on error, undefined variables, and print commands
#set -eux
set -e

# Default values (can be overridden by command-line arguments)
SESSION_NAME="${1:-imagegen}"
PORT="${2:-8080}"
LISTEN_ADDRESS="${3:-127.0.0.1}"
LOWVRAM="${4:-false}"

if tmux has-session -t "$SESSION_NAME" 2>/dev/null; then
    echo "Session '$SESSION_NAME' already exists. No action taken."
    exit 0
fi

# Activate virtual environment
source .venv/bin/activate

# Check if session already exists
if ! tmux has-session -t "$SESSION_NAME" 2>/dev/null; then
    # Create new detached session
    tmux new-session -d -s "$SESSION_NAME"

    # Prepare the command
    BASE_COMMAND="python main.py --port $PORT --listen $LISTEN_ADDRESS"
    if [ "$LOWVRAM" = "true" ]; then
        FULL_COMMAND="$BASE_COMMAND --lowvram --preview-method auto --use-split-cross-attention"
    else
        FULL_COMMAND="$BASE_COMMAND"
    fi

    # Start the ComfyUI application
    tmux send-keys -t "$SESSION_NAME" "$FULL_COMMAND" C-m

    echo "Session '$SESSION_NAME' created and ComfyUI started."
    echo "To connect, type: tmux attach -t $SESSION_NAME"
else
    echo "Session '$SESSION_NAME' already exists. Exiting."
    exit 1
fi