#!/usr/bin/env bash
set -euo pipefail

# split-wezterm — Split the current WezTerm pane into a rows x cols grid

DRYRUN=false

usage() {
    cat <<'EOF'
Usage: split-wezterm [--dryrun|-n] <rows> <cols> [--] [command...]

Split the current WezTerm pane into an evenly-sized grid.

Arguments:
  rows        Number of rows (1-10)
  cols        Number of columns (1-10)
  command...  Optional command to run in each new pane

Flags:
  --dryrun, -n   Preview commands without executing
  --help, -h     Show this help

Examples:
  split-wezterm 3 4
  split-wezterm 2 2 -- htop
  split-wezterm --dryrun 3 4 -- claude
EOF
}

die() { echo "Error: $*" >&2; exit 1; }

# Parse flags
while [[ $# -gt 0 ]]; do
    case "$1" in
        --help|-h) usage; exit 0 ;;
        --dryrun|-n) DRYRUN=true; shift ;;
        --) shift; break ;;
        -*) die "Unknown flag: $1" ;;
        *) break ;;
    esac
done

[[ $# -lt 2 ]] && { usage >&2; exit 1; }

ROWS="$1"; shift
COLS="$1"; shift
# Skip -- separator if present
[[ "${1:-}" == "--" ]] && shift
CMD=("$@")

# Validate dimensions
[[ "$ROWS" =~ ^[0-9]+$ ]] || die "rows must be a positive integer, got: $ROWS"
[[ "$COLS" =~ ^[0-9]+$ ]] || die "cols must be a positive integer, got: $COLS"
(( ROWS >= 1 && ROWS <= 10 )) || die "rows must be 1-10, got: $ROWS"
(( COLS >= 1 && COLS <= 10 )) || die "cols must be 1-10, got: $COLS"

if (( ROWS == 1 && COLS == 1 )); then
    echo "Already a single pane, nothing to do."
    exit 0
fi

# Detect WezTerm CLI
detect_cli() {
    if command -v wezterm &>/dev/null; then
        echo "wezterm"
    elif flatpak info org.wezfurlong.wezterm &>/dev/null 2>&1; then
        echo "flatpak run org.wezfurlong.wezterm"
    else
        die "WezTerm CLI not found. Install WezTerm or ensure it's in PATH."
    fi
}

WEZTERM_CLI=$(detect_cli)

# Get the starting pane ID
if [[ -n "${WEZTERM_PANE:-}" ]]; then
    START_PANE="$WEZTERM_PANE"
else
    # Try to find the active pane from wezterm cli list
    START_PANE=$($WEZTERM_CLI cli list --format json 2>/dev/null \
        | python3 -c "import sys,json; panes=json.load(sys.stdin); print(panes[0]['pane_id'])" 2>/dev/null) \
        || die "Cannot determine current pane. Set WEZTERM_PANE or run from inside WezTerm."
fi

# Run or preview a wezterm cli command; capture stdout (the new pane ID)
run_cli() {
    if $DRYRUN; then
        echo "[dryrun] $WEZTERM_CLI cli $*" >&2
        echo "dry-$$-$RANDOM"  # fake pane ID for dryrun tracking
    else
        $WEZTERM_CLI cli "$@"
    fi
}

# Split a pane into N equal parts along an axis.
# Usage: split_axis <pane_id> <count> <direction_flag> <result_array_name>
# direction_flag: --bottom (rows) or --horizontal (cols)
# Populates the named array with all pane IDs (original + new)
split_axis() {
    local pane_id="$1" count="$2" direction="$3" result_var="$4"
    local -a pane_ids=("$pane_id")

    local i percent new_pane
    for (( i = 0; i < count - 1; i++ )); do
        percent=$(( (count - 1 - i) * 100 / (count - i) ))
        local cmd_args=(split-pane --pane-id "$pane_id" "$direction" --percent "$percent")
        if [[ ${#CMD[@]} -gt 0 ]]; then
            cmd_args+=(-- "${CMD[@]}")
        fi
        new_pane=$(run_cli "${cmd_args[@]}")
        pane_id="$new_pane"
        pane_ids+=("$new_pane")
    done

    # Copy results to the named variable
    eval "$result_var=(\"\${pane_ids[@]}\")"
}

# Step 1: Split into rows
split_axis "$START_PANE" "$ROWS" "--bottom" row_panes

# Step 2: Split each row into columns
for row_pane in "${row_panes[@]}"; do
    if (( COLS > 1 )); then
        split_axis "$row_pane" "$COLS" "--horizontal" _col_panes
    fi
done

# Step 3: If a command was given, run it in the original pane too
if [[ ${#CMD[@]} -gt 0 ]]; then
    cmd_text="${CMD[*]}"
    if $DRYRUN; then
        echo "[dryrun] $WEZTERM_CLI cli send-text --pane-id $START_PANE --no-paste -- \"$cmd_text\"" >&2
    else
        # Send the command text followed by enter
        $WEZTERM_CLI cli send-text --pane-id "$START_PANE" --no-paste -- "$cmd_text
"
    fi
fi

TOTAL=$(( ROWS * COLS ))
if $DRYRUN; then
    echo "[dryrun] Would create ${ROWS}x${COLS} grid ($TOTAL panes)"
else
    echo "Created ${ROWS}x${COLS} grid ($TOTAL panes)"
fi