#!/usr/bin/env bash
# Tmux Session Manager - Manage multiple named tmux sessions with persistence
# Usage: tmux-session-manager.sh [session-name] [action]
# Actions: save, restore, list, delete

set -euo pipefail

SESSION_DIR="${XDG_DATA_HOME:-$HOME/.local/share}/tmux/sessions"
RESURRECT_DIR="${XDG_DATA_HOME:-$HOME/.local/share}/tmux/resurrect"

# Create directories if they don't exist
mkdir -p "$SESSION_DIR" "$RESURRECT_DIR"

show_usage() {
    cat <<EOF
Tmux Session Manager - Manage multiple named tmux sessions with persistence

Usage: tmux-session-manager [command] [options]

Commands:
  list                     List all saved session configurations
  create <name>           Create or switch to a named session
  save <name>             Manually save current session state
  restore <name>          Restore a specific named session
  delete <name>           Delete a named session configuration
  restore-all             Restore all sessions at once
  auto-save               Automatically save all active sessions (for cron/timer)

Examples:
  tmux-session-manager create main       # Create/switch to 'main' session
  tmux-session-manager save work         # Save current session as 'work'
  tmux-session-manager restore dev       # Restore 'dev' session
  tmux-session-manager list              # List all sessions
EOF
}

last_snapshot_sessions() {
    local last_file="$RESURRECT_DIR/last"

    if [ ! -f "$last_file" ]; then
        return 0
    fi

    awk -F '\t' '$1 == "pane" { print $2 }' "$last_file" | sort -u
}

session_path_matches_name() {
    local session="$1"
    local desired_name="$2"
    local desired_basename="$desired_name"
    local path=""

    if [[ "$desired_name" == _* ]]; then
        desired_basename=".${desired_name#_}"
    fi

    while IFS= read -r path; do
        [ -n "$path" ] || continue
        if [ "$(basename "$path")" = "$desired_basename" ]; then
            return 0
        fi
    done < <(tmux list-panes -t "=${session}" -F "#{pane_current_path}" 2>/dev/null)

    if [ "$desired_name" = "main" ]; then
        while IFS= read -r path; do
            [ -n "$path" ] || continue
            if [ "$path" = "$HOME" ]; then
                return 0
            fi
        done < <(tmux list-panes -t "=${session}" -F "#{pane_current_path}" 2>/dev/null)
    fi

    return 1
}

run_resurrect_restore() {
    local restore_script
    restore_script="$(tmux show-options -gv @resurrect-restore-script-path 2>/dev/null || true)"

    if [ -z "$restore_script" ]; then
        echo "Error: Could not determine resurrect restore script path"
        return 1
    fi

    tmux run-shell "$restore_script"
}

run_resurrect_save() {
    local save_script
    save_script="$(tmux show-options -gv @resurrect-save-script-path 2>/dev/null || true)"

    if [ -z "$save_script" ]; then
        echo "Error: Could not determine resurrect save script path"
        return 1
    fi

    tmux run-shell "$save_script quiet"
}

save_session() {
    local session_name="$1"
    
    # Get the current session if not specified
    if [ -z "$session_name" ]; then
        session_name=$(tmux display-message -p '#{session_name}' 2>/dev/null || echo "")
        if [ -z "$session_name" ]; then
            echo "Error: No active tmux session found"
            return 1
        fi
    fi
    
    # Check if session exists
    if ! tmux has-session -t "$session_name" 2>/dev/null; then
        echo "Error: Session '$session_name' does not exist"
        return 1
    fi
    
    # Save session using resurrect script
    run_resurrect_save
    
    echo "Saving session '$session_name'..."
    sleep 1
    
    # Also create a backup of the session state
    local backup_file="$SESSION_DIR/${session_name}.state"
    tmux capture-pane -t "$session_name" -p > "$backup_file" || true
    
    echo "Session '$session_name' saved"
}

restore_session() {
    local session_name="$1"
    local bootstrap_session="__tsm_restore_bootstrap"
    local trigger_session=""
    local restored=0
    local renamed=0
    local candidate_session=""
    local -a before_sessions=()
    local -a after_sessions=()
    local -a new_sessions=()
    
    if [ -z "$session_name" ]; then
        echo "Error: Session name required"
        show_usage
        return 1
    fi
    
    # Use resurrect to restore the most recent snapshot.
    # A target session might not currently exist, so trigger restore from any live session.
    echo "Restoring session '$session_name' from latest resurrect save..."

    while IFS= read -r session; do
        [ -n "$session" ] && before_sessions+=("$session")
    done < <(tmux list-sessions -F "#{session_name}" 2>/dev/null || true)

    if ! tmux list-sessions >/dev/null 2>&1; then
        tmux new-session -d -s "$bootstrap_session"
        trigger_session="$bootstrap_session"
    else
        while IFS= read -r session; do
            if [ -n "$session" ]; then
                trigger_session="$session"
                break
            fi
        done < <(tmux list-sessions -F "#{session_name}" 2>/dev/null)
    fi

    if [ -z "$trigger_session" ]; then
        echo "Error: Could not find a tmux session to trigger restore"
        return 1
    fi

    run_resurrect_restore

    for _ in $(seq 1 20); do
        if tmux has-session -t "$session_name" 2>/dev/null; then
            restored=1
            break
        fi
        sleep 0.25
    done

    while IFS= read -r session; do
        [ -n "$session" ] && after_sessions+=("$session")
    done < <(tmux list-sessions -F "#{session_name}" 2>/dev/null || true)

    if [ "$restored" -eq 0 ]; then
        for session in "${after_sessions[@]}"; do
            local existed=0
            for before in "${before_sessions[@]}"; do
                if [ "$session" = "$before" ]; then
                    existed=1
                    break
                fi
            done
            if [ "$existed" -eq 0 ]; then
                new_sessions+=("$session")
            fi
        done

        # Give resurrect a moment to fully materialize pane paths before matching.
        for _ in $(seq 1 20); do
            candidate_session=""

            for session in "${new_sessions[@]}"; do
                if session_path_matches_name "$session" "$session_name"; then
                    candidate_session="$session"
                    break
                fi
            done

            # If no brand-new session matched, fall back to any current session.
            # This handles cases where restore reused existing numeric sessions.
            if [ -z "$candidate_session" ]; then
                for session in "${after_sessions[@]}"; do
                    if [ "$session" = "$session_name" ]; then
                        continue
                    fi
                    if session_path_matches_name "$session" "$session_name"; then
                        candidate_session="$session"
                        break
                    fi
                done
            fi

            if [ -n "$candidate_session" ]; then
                break
            fi

            sleep 0.25
        done

        if [ -z "$candidate_session" ] && [ "${#new_sessions[@]}" -eq 1 ]; then
            candidate_session="${new_sessions[0]}"
        fi

        if [ -n "$candidate_session" ] && ! tmux has-session -t "$session_name" 2>/dev/null; then
            tmux rename-session -t "=${candidate_session}" "$session_name"
            restored=1
            renamed=1
        fi
    fi

    if [ "$trigger_session" = "$bootstrap_session" ] && tmux has-session -t "$bootstrap_session" 2>/dev/null; then
        tmux kill-session -t "$bootstrap_session" 2>/dev/null || true
    fi

    if [ "$restored" -eq 1 ]; then
        if [ "$renamed" -eq 1 ]; then
            echo "Session '$session_name' restored by renaming '$candidate_session'"
        else
            echo "Session '$session_name' restored and available"
        fi
    else
        local snapshot_sessions
        snapshot_sessions="$(last_snapshot_sessions | tr '\n' ' ' | sed 's/[[:space:]]*$//')"
        if [ -n "$snapshot_sessions" ]; then
            echo "Restore was triggered, but '$session_name' was not found in the latest snapshot (contains: $snapshot_sessions)"
        else
            echo "Restore was triggered, but '$session_name' was not found and no snapshot sessions were detected"
        fi
        return 1
    fi
}

create_session() {
    local session_name="$1"
    
    if [ -z "$session_name" ]; then
        echo "Error: Session name required"
        show_usage
        return 1
    fi
    
    # Check if session already exists
    if tmux has-session -t "$session_name" 2>/dev/null; then
        echo "Session '$session_name' already exists, attaching..."
        tmux attach-session -t "=${session_name}"
    else
        echo "Creating new session '$session_name'..."
        tmux new-session -d -s "$session_name"
        tmux send-keys -t "$session_name" "# Session: $session_name" Enter
        echo "Session '$session_name' created"
        tmux attach-session -t "=${session_name}"
    fi
}

list_sessions() {
    echo "Available tmux sessions:"
    tmux list-sessions -F "#{session_name}: #{session_windows} windows" 2>/dev/null || echo "No tmux sessions found"

    echo ""
    echo "Session names in latest resurrect snapshot:"
    local snapshot_sessions
    snapshot_sessions="$(last_snapshot_sessions)"
    if [ -n "$snapshot_sessions" ]; then
        echo "$snapshot_sessions"
    else
        echo "No resurrect snapshot found"
    fi
    
    echo ""
    echo "Saved session labels (local backups):"
    if [ -d "$SESSION_DIR" ]; then
        ls -1 "$SESSION_DIR" 2>/dev/null | sed 's/.state$//' | sort | uniq || echo "No saved sessions"
    else
        echo "No saved sessions"
    fi
}

delete_session() {
    local session_name="$1"
    
    if [ -z "$session_name" ]; then
        echo "Error: Session name required"
        return 1
    fi
    
    # Kill the tmux session if it exists
    if tmux has-session -t "$session_name" 2>/dev/null; then
        tmux kill-session -t "=${session_name}"
        echo "Killed tmux session '$session_name'"
    fi
    
    # Remove saved state files
    rm -f "$SESSION_DIR/${session_name}.state"
    echo "Deleted session configuration '$session_name'"
}

restore_all_sessions() {
    local bootstrap_session="__tsm_restore_bootstrap"
    local trigger_session=""

    echo "Restoring all sessions from latest resurrect save..."

    if ! tmux list-sessions >/dev/null 2>&1; then
        tmux new-session -d -s "$bootstrap_session"
        trigger_session="$bootstrap_session"
    else
        while IFS= read -r session; do
            if [ -n "$session" ]; then
                trigger_session="$session"
                break
            fi
        done < <(tmux list-sessions -F "#{session_name}" 2>/dev/null)
    fi

    if [ -z "$trigger_session" ]; then
        echo "Error: Could not find a tmux session to trigger restore"
        return 1
    fi

    run_resurrect_restore
    sleep 1

    if [ "$trigger_session" = "$bootstrap_session" ] && tmux has-session -t "$bootstrap_session" 2>/dev/null; then
        tmux kill-session -t "=${bootstrap_session}" 2>/dev/null || true
    fi

    echo "Restore triggered for all sessions"
}

auto_save() {
    echo "Auto-saving all active tmux sessions..."
    while IFS= read -r session; do
        save_session "$session" || true
    done < <(tmux list-sessions -F "#{session_name}" 2>/dev/null)
    echo "All sessions auto-saved"
}

# Main command dispatcher
main() {
    local command="${1:-list}"
    local arg="${2:-}"
    
    case "$command" in
        list)
            list_sessions
            ;;
        create)
            create_session "$arg"
            ;;
        save)
            save_session "$arg"
            ;;
        restore)
            restore_session "$arg"
            ;;
        delete)
            delete_session "$arg"
            ;;
        restore-all)
            restore_all_sessions
            ;;
        auto-save)
            auto_save
            ;;
        --help|-h|help)
            show_usage
            ;;
        *)
            echo "Unknown command: $command"
            show_usage
            return 1
            ;;
    esac
}

main "$@"