#!/bin/bash

# Launch Claude Code instances in tmux grid
# Usage: claude-grid [--yolo] [panes] [dir1] [dir2] ... [dirN]
#   --yolo: pass --dangerously-skip-permissions to claude
#   panes:  2, 3, 4, 6 (or any supported count)
#   dirs:   optional per-pane directories. If fewer dirs than panes, remaining use current dir.
#           If only one dir given, all panes use it. If none, all use current dir.
#
# Examples:
#   claude-grid 4                                      # 4 panes, all in current dir
#   claude-grid 3 ~/project                            # 3 panes, all in ~/project
#   claude-grid --yolo 6 dir1 dir2 dir3 dir4 dir5 dir6 # 6 panes, skip permissions

CLAUDE_CMD="claude"

# Check for --yolo flag
if [ "$1" = "--yolo" ]; then
    CLAUDE_CMD="claude --dangerously-skip-permissions"
    shift
fi

PANES="${1:-4}"
shift

# Collect directories (remaining args)
DIRS=()
DEFAULT_DIR="$(pwd)"
for arg in "$@"; do
    DIRS+=("$(cd "$arg" 2>/dev/null && pwd || echo "$arg")")
done

# Helper: get directory for pane N
get_dir() {
    local idx=$1
    if [ ${#DIRS[@]} -eq 0 ]; then
        echo "$DEFAULT_DIR"
    elif [ ${#DIRS[@]} -eq 1 ]; then
        echo "${DIRS[0]}"
    elif [ $idx -lt ${#DIRS[@]} ]; then
        echo "${DIRS[$idx]}"
    else
        echo "$DEFAULT_DIR"
    fi
}

SESSION="claude-grid-$$"

# Create new tmux session with first pane
tmux new-session -d -s "$SESSION" -c "$(get_dir 0)"

case "$PANES" in
    2)
        # 2 panes side by side
        tmux split-window -h -t "$SESSION" -c "$(get_dir 1)"
        ;;
    3)
        # 3 panes: 2 top + 1 bottom
        tmux split-window -h -t "$SESSION" -c "$(get_dir 1)"
        tmux split-window -v -t "$SESSION" -c "$(get_dir 2)"
        tmux select-layout -t "$SESSION" main-horizontal
        ;;
    4)
        # 4 panes 2x2
        tmux split-window -h -t "$SESSION" -c "$(get_dir 1)"
        tmux split-window -v -t "$SESSION" -c "$(get_dir 2)"
        tmux select-pane -t "$SESSION" -L
        tmux split-window -v -t "$SESSION" -c "$(get_dir 3)"
        tmux select-layout -t "$SESSION" tiled
        ;;
    6)
        # 6 panes: 3 rows x 2 cols using pane IDs (not indices)
        #   0 | 1
        #   2 | 3
        #   4 | 5

        # Get pane ID of the first pane
        P0=$(tmux list-panes -t "$SESSION" -F '#{pane_id}')

        # Split right -> 2 columns
        P1=$(tmux split-window -h -t "$P0" -c "$(get_dir 1)" -P -F '#{pane_id}')

        # Split left column into 3 rows
        P2=$(tmux split-window -v -t "$P0" -c "$(get_dir 2)" -P -F '#{pane_id}')
        P4=$(tmux split-window -v -t "$P2" -c "$(get_dir 4)" -P -F '#{pane_id}')

        # Split right column into 3 rows
        P3=$(tmux split-window -v -t "$P1" -c "$(get_dir 3)" -P -F '#{pane_id}')
        P5=$(tmux split-window -v -t "$P3" -c "$(get_dir 5)" -P -F '#{pane_id}')
        ;;
    *)
        echo "Usage: claude-grid [--yolo] [2|3|4|6] [dir1] [dir2] ... [dirN]"
        echo ""
        echo "Options:"
        echo "  --yolo    Pass --dangerously-skip-permissions to claude"
        echo ""
        echo "Layouts:"
        echo "  2 - side by side"
        echo "  3 - 2 top + 1 bottom"
        echo "  4 - 2x2 grid"
        echo "  6 - 3x2 grid (3 top + 3 bottom)"
        echo ""
        echo "Directories are optional. Pass one per pane, or one for all, or none for cwd."
        tmux kill-session -t "$SESSION" 2>/dev/null
        exit 1
        ;;
esac

# Send claude to all panes
for pane in $(tmux list-panes -t "$SESSION" -F '#{pane_id}'); do
    tmux send-keys -t "$pane" "$CLAUDE_CMD" C-m
done

# If --yolo, auto-accept the bypass permissions disclaimer
# The disclaimer renders a select widget: cursor starts on "1. No, exit"
# We need to press Down to move to "2. Yes, I accept", then Enter to confirm
if [ "$CLAUDE_CMD" != "claude" ]; then
    sleep 6  # wait for all panes to render the disclaimer
    for pane in $(tmux list-panes -t "$SESSION" -F '#{pane_id}'); do
        tmux send-keys -t "$pane" Down
        sleep 0.3
        tmux send-keys -t "$pane" Enter
        sleep 1  # wait for claude to fully boot before next pane
    done
fi

tmux select-pane -t "${SESSION}:0.0"
tmux attach-session -t "$SESSION"
