#!/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: playground)
#   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 (uses system hostname):
#      ./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 -e

# Default values (can be overridden by command-line arguments)
# Use "playground" as default session name if not provided
SESSION_NAME="${1:-playground}"
OPEN_WEBUI_PORT="${2:-8080}"
LITELLM_PORT="${3:-8031}" # consistent with litten
COMFYUI_PORT="${4:-8081}"
COMFYUI_LOWVRAM="${5:-false}"

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

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

    # Create new detached tmux session, starting with open-webui and litellm
    tmux new-session -d -s "$SESSION_NAME" -n "textgen"
    tmux send-keys -t "$SESSION_NAME":0 "cd /opt/playground && source .venv_open_webui/bin/activate.fish && open-webui serve --port $OPEN_WEBUI_PORT" C-m
    tmux split-window -v -t "$SESSION_NAME":0
    tmux send-keys -t "$SESSION_NAME":0.1 "cd /opt/playground && source .venv_litellm/bin/activate.fish && source .env && litellm --telemetry False --config ./litellm.yaml --host 127.0.0.1 --port $LITELLM_PORT" C-m

    # Start ollama in a new window
    tmux new-window -t "$SESSION_NAME" -n "ollama"
    tmux send-keys -t "$SESSION_NAME":1 "cd /opt/ollama && ./bin/ollama serve" C-m

    # Prepare the ComfyUI launch command
    COMFYUI_BASE_COMMAND="python main.py --port $COMFYUI_PORT --listen 127.0.0.1"
    if [ "$COMFYUI_LOWVRAM" = "true" ]; then
        COMFYUI_LAUNCH_COMMAND="cd /opt/comfyui && source .venv/bin/activate.fish && $COMFYUI_BASE_COMMAND --lowvram --preview-method auto --use-split-cross-attention"
    else
        COMFYUI_LAUNCH_COMMAND="cd /opt/comfyui && source .venv/bin/activate.fish && $COMFYUI_BASE_COMMAND"
    fi

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

    # Start ComfyUI in window 2
    tmux new-window -t "$SESSION_NAME" -n "comfyui"
    tmux send-keys -t "$SESSION_NAME":2 "$COMFYUI_LAUNCH_COMMAND" C-m

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