#!/bin/bash
# comfygit-dev - Development launcher for ComfyGit
#
# Usage: comfygit-dev [-s|--simulator] [/path/to/ComfyUI] [-- additional args]
#
# Options:
#   -s, --simulator    Use local Docker simulator instead of RunPod API
#
# Sets up dev environment and launches ComfyUI with comfygit-manager

set -e

# Resolve symlinks to get actual script location
SCRIPT_PATH="$(readlink -f "${BASH_SOURCE[0]}")"
SCRIPT_DIR="$(dirname "$SCRIPT_PATH")"
MANAGER_PATH="$(dirname "$SCRIPT_DIR")"
CORE_PATH="/home/akatzfey/projects/comfyhub/comfygit/packages/core"
DEFAULT_COMFYUI="/home/akatzfey/projects/ComfyUI_dev/ComfyUI"
RUNPOD_IMAGE="runpod/pytorch:2.1.0-py3.10-cuda11.8.0-devel-ubuntu22.04"

# Parse arguments
COMFYUI_PATH="$DEFAULT_COMFYUI"
EXTRA_ARGS=()
SIMULATOR_MODE=false

while [[ $# -gt 0 ]]; do
    case $1 in
        -s|--simulator)
            SIMULATOR_MODE=true
            shift
            ;;
        --)
            shift
            EXTRA_ARGS=("$@")
            break
            ;;
        *)
            if [[ -d "$1" ]]; then
                COMFYUI_PATH="$1"
            else
                EXTRA_ARGS+=("$1")
            fi
            shift
            ;;
    esac
done

# Validate paths
if [[ ! -d "$COMFYUI_PATH" ]]; then
    echo "Error: ComfyUI not found at $COMFYUI_PATH"
    exit 1
fi

if [[ ! -f "$COMFYUI_PATH/main.py" ]]; then
    echo "Error: $COMFYUI_PATH doesn't look like a ComfyUI installation (no main.py)"
    exit 1
fi

echo "[comfygit-dev] ComfyUI: $COMFYUI_PATH"
echo "[comfygit-dev] Manager: $MANAGER_PATH"
echo "[comfygit-dev] Core:    $CORE_PATH"

# Check for existing processes on port 8188
check_and_kill_zombies() {
    local PORT=8188
    local PID=$(lsof -t -i :$PORT 2>/dev/null | head -1)

    if [[ -n "$PID" ]]; then
        local CMD=$(ps -p "$PID" -o cmd= 2>/dev/null || echo "unknown")
        echo "[comfygit-dev] Port $PORT is in use by PID $PID"
        echo "[comfygit-dev] Command: $CMD"

        read -p "[comfygit-dev] Kill existing process? [y/N] " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            # Also kill parent if it's an orchestrator
            local PARENT_PID=$(ps -p "$PID" -o ppid= 2>/dev/null | tr -d ' ')
            if [[ -n "$PARENT_PID" ]] && [[ "$PARENT_PID" != "1" ]]; then
                local PCMD=$(ps -p "$PARENT_PID" -o cmd= 2>/dev/null || echo "")
                if [[ "$PCMD" == *"orchestrator.py"* ]]; then
                    echo "[comfygit-dev] Killing orchestrator (PID $PARENT_PID)..."
                    kill "$PARENT_PID" 2>/dev/null || true
                fi
            fi
            echo "[comfygit-dev] Killing process (PID $PID)..."
            kill "$PID" 2>/dev/null || true
            sleep 1
        else
            echo "[comfygit-dev] Aborting. Kill the process manually or use a different port."
            exit 1
        fi
    fi
}

check_and_kill_zombies

# Ensure custom_nodes directory exists
mkdir -p "$COMFYUI_PATH/custom_nodes"

# Check/create comfygit-manager symlink
MANAGER_LINK="$COMFYUI_PATH/custom_nodes/comfygit-manager"
if [[ -L "$MANAGER_LINK" ]]; then
    CURRENT_TARGET="$(readlink "$MANAGER_LINK")"
    if [[ "$CURRENT_TARGET" != "$MANAGER_PATH" ]]; then
        echo "[comfygit-dev] Updating manager symlink (was: $CURRENT_TARGET)"
        rm "$MANAGER_LINK"
        ln -s "$MANAGER_PATH" "$MANAGER_LINK"
    else
        echo "[comfygit-dev] Manager symlink OK"
    fi
elif [[ -e "$MANAGER_LINK" ]]; then
    echo "Error: $MANAGER_LINK exists but is not a symlink"
    exit 1
else
    echo "[comfygit-dev] Creating manager symlink"
    ln -s "$MANAGER_PATH" "$MANAGER_LINK"
fi

# Check/create .venv with comfygit-core
VENV_PATH="$COMFYUI_PATH/.venv"
if [[ ! -d "$VENV_PATH" ]]; then
    echo "[comfygit-dev] Creating venv..."
    cd "$COMFYUI_PATH"
    uv venv
    uv pip install -r requirements.txt
    uv pip install -e "$CORE_PATH"
else
    # Check if comfygit-core is installed as editable
    INSTALLED=$(cd "$COMFYUI_PATH" && uv pip show comfygit-core 2>/dev/null | grep "Editable project location" || true)
    if [[ -z "$INSTALLED" ]]; then
        echo "[comfygit-dev] Installing comfygit-core as editable..."
        cd "$COMFYUI_PATH"
        uv pip install -e "$CORE_PATH"
    elif [[ "$INSTALLED" != *"$CORE_PATH"* ]]; then
        echo "[comfygit-dev] Reinstalling comfygit-core (wrong path)..."
        cd "$COMFYUI_PATH"
        uv pip install -e "$CORE_PATH"
    else
        echo "[comfygit-dev] Core editable install OK"
    fi
fi

# Clean stale orchestrator venv if needed
ORCH_VENV="$MANAGER_PATH/server/.orchestrator_venv"
if [[ -d "$ORCH_VENV" ]]; then
    # Check if it has the right core installed
    ORCH_CORE=$("$ORCH_VENV/bin/pip" show comfygit-core 2>/dev/null | grep "Editable project location" || true)
    if [[ -z "$ORCH_CORE" ]] || [[ "$ORCH_CORE" != *"$CORE_PATH"* ]]; then
        echo "[comfygit-dev] Removing stale orchestrator venv..."
        rm -rf "$ORCH_VENV"
    fi
fi

# Update core library source in all managed environments' pyproject.toml
# This ensures uv sync uses the dev version of comfygit-core
WORKSPACE_PATH="$HOME/comfygit"
if [[ -d "$WORKSPACE_PATH/environments" ]]; then
    for ENV_DIR in "$WORKSPACE_PATH/environments"/*/; do
        ENV_NAME=$(basename "$ENV_DIR")
        PYPROJECT="$ENV_DIR.cec/pyproject.toml"

        if [[ -f "$PYPROJECT" ]]; then
            # Check if comfygit-core source is already set to dev path
            if ! grep -q "tool.uv.sources.comfygit-core" "$PYPROJECT" 2>/dev/null; then
                echo "[comfygit-dev] Adding dev core source to: $ENV_NAME"
                # Append the source configuration
                cat >> "$PYPROJECT" << EOF

[tool.uv.sources.comfygit-core]
path = "$CORE_PATH"
editable = true
EOF
            elif ! grep -A2 "tool.uv.sources.comfygit-core" "$PYPROJECT" | grep -q "$CORE_PATH"; then
                echo "[comfygit-dev] Updating dev core source in: $ENV_NAME"
                # Use Python to update the TOML properly
                python3 << PYEOF
import re
with open("$PYPROJECT", "r") as f:
    content = f.read()
# Remove existing comfygit-core source section
content = re.sub(r'\[tool\.uv\.sources\.comfygit-core\][^\[]*', '', content)
# Add new section
content = content.rstrip() + '''

[tool.uv.sources.comfygit-core]
path = "$CORE_PATH"
editable = true
'''
with open("$PYPROJECT", "w") as f:
    f.write(content)
PYEOF
            fi
        fi
    done
    echo "[comfygit-dev] Managed environments pyproject.toml OK"
fi

# Refresh system-nodes dependency group from comfygit-manager/pyproject.toml
# This ensures new dependencies (like zeroconf) are picked up
if [[ -d "$WORKSPACE_PATH/environments" ]] && [[ -f "$MANAGER_PATH/pyproject.toml" ]]; then
    echo "[comfygit-dev] Refreshing system-nodes dependencies from manager..."
    for ENV_DIR in "$WORKSPACE_PATH/environments"/*/; do
        ENV_NAME=$(basename "$ENV_DIR")
        PYPROJECT="$ENV_DIR.cec/pyproject.toml"

        if [[ -f "$PYPROJECT" ]]; then
            # Use uv run from manager dir to get tomlkit
            cd "$MANAGER_PATH"
            uv run python << PYEOF
import tomlkit

# Read manager dependencies
with open("$MANAGER_PATH/pyproject.toml", "r") as f:
    manager_doc = tomlkit.load(f)
manager_deps = manager_doc.get("project", {}).get("dependencies", [])

if not manager_deps:
    exit(0)

# Read environment pyproject
with open("$PYPROJECT", "r") as f:
    env_doc = tomlkit.load(f)

# Get current system-nodes deps
if "dependency-groups" not in env_doc:
    env_doc["dependency-groups"] = {}
current_deps = set(env_doc.get("dependency-groups", {}).get("system-nodes", []))
new_deps = set(manager_deps)

# Only update if different
if current_deps != new_deps:
    print(f"[comfygit-dev] Updating system-nodes in $ENV_NAME: {sorted(new_deps)}")
    env_doc["dependency-groups"]["system-nodes"] = sorted(list(new_deps))
    with open("$PYPROJECT", "w") as f:
        tomlkit.dump(env_doc, f)
PYEOF
        fi
    done
fi

# Add docker dependency to managed environments when simulator mode is enabled
if [[ "$SIMULATOR_MODE" == "true" ]] && [[ -d "$WORKSPACE_PATH/environments" ]]; then
    for ENV_DIR in "$WORKSPACE_PATH/environments"/*/; do
        ENV_NAME=$(basename "$ENV_DIR")
        PYPROJECT="$ENV_DIR.cec/pyproject.toml"

        if [[ -f "$PYPROJECT" ]]; then
            # Check if docker is already in dependencies
            if ! grep -q '"docker' "$PYPROJECT" 2>/dev/null; then
                echo "[comfygit-dev] Adding docker dependency to: $ENV_NAME"
                # Use tomlkit for safe TOML manipulation
                python3 << PYEOF
import tomlkit

with open("$PYPROJECT", "r") as f:
    doc = tomlkit.load(f)

deps = doc.get("project", {}).get("dependencies", [])
if not any("docker" in d for d in deps):
    # Add docker in alphabetical position
    deps.append("docker>=7.1.0")
    # Sort to maintain order
    deps.sort(key=lambda x: x.lower())
    doc["project"]["dependencies"] = deps

    with open("$PYPROJECT", "w") as f:
        tomlkit.dump(doc, f)
PYEOF
            fi
        fi
    done
    echo "[comfygit-dev] Simulator docker dependency added to managed environments"
fi

# Export dev mode environment variables
export COMFYGIT_DEV_CORE_PATH="$CORE_PATH"
export COMFYGIT_DEV_MANAGER_PATH="$MANAGER_PATH"

# Setup simulator mode if requested
if [[ "$SIMULATOR_MODE" == "true" ]]; then
    echo ""
    echo "[comfygit-dev] ========== SIMULATOR MODE =========="

    # Check if RunPod image exists, prompt to pull if not
    if ! docker images "$RUNPOD_IMAGE" -q 2>/dev/null | grep -q .; then
        echo "[comfygit-dev] RunPod base image not found locally."
        read -p "[comfygit-dev] Pull image now? (~15GB) [y/N] " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            echo "[comfygit-dev] Pulling image (this may take a while)..."
            docker pull "$RUNPOD_IMAGE"
        else
            echo "[comfygit-dev] Image will be pulled on first deploy (may cause delay)"
        fi
    fi

    # Set simulator environment
    export DEPLOY_BACKEND=local

    # Detect GPU availability
    if command -v nvidia-smi &> /dev/null && nvidia-smi &> /dev/null; then
        export SIMULATOR_GPU_MODE=${SIMULATOR_GPU_MODE:-gpu}
        echo "[comfygit-dev] GPU detected - using GPU mode"
    else
        export SIMULATOR_GPU_MODE=cpu
        echo "[comfygit-dev] No GPU detected - using CPU mode"
    fi

    echo "[comfygit-dev] DEPLOY_BACKEND=$DEPLOY_BACKEND"
    echo "[comfygit-dev] SIMULATOR_GPU_MODE=$SIMULATOR_GPU_MODE"
    echo "[comfygit-dev] Deploy operations will use local Docker instead of RunPod."
    echo "[comfygit-dev] ======================================"
    echo ""
fi

echo "[comfygit-dev] Starting ComfyUI..."
cd "$COMFYUI_PATH"
exec uv run main.py "${EXTRA_ARGS[@]}"
