#!/bin/bash
#
# create_worktrees.sh - Set up git worktrees and tmux sessions for multi-agent orchestration
#
# This script creates isolated workspaces for multiple AI agents using:
# - Git worktrees: Separate working directories per agent
# - tmux sessions: Process isolation and terminal management
#
# Usage:
#   ./create_worktrees.sh [options]
#
# Options:
#   -p, --project PATH   Project root directory (default: current directory)
#   -c, --config FILE    Config file with agent definitions (optional)
#   -b, --base BRANCH    Base branch to create worktrees from (default: main)
#   -n, --dry-run        Show what would be done without executing
#   -h, --help           Show this help message
#
# Example:
#   ./create_worktrees.sh -p /path/to/repo -b main
#
# The script will create:
#   - agent-claude-code/     Worktree for Claude Code agent
#   - agent-gemini/          Worktree for Gemini CLI agent
#   - agent-codex/           Worktree for Codex CLI agent
#

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Default values
PROJECT_ROOT="$(pwd)"
BASE_BRANCH="main"
DRY_RUN=false
CONFIG_FILE=""

# Default agents configuration
declare -A AGENTS
AGENTS["claude-code"]="feature/claude-code"
AGENTS["gemini"]="feature/gemini"
AGENTS["codex"]="feature/codex"

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -p|--project)
            PROJECT_ROOT="$2"
            shift 2
            ;;
        -c|--config)
            CONFIG_FILE="$2"
            shift 2
            ;;
        -b|--base)
            BASE_BRANCH="$2"
            shift 2
            ;;
        -n|--dry-run)
            DRY_RUN=true
            shift
            ;;
        -h|--help)
            head -40 "$0" | tail -35
            exit 0
            ;;
        *)
            echo -e "${RED}Unknown option: $1${NC}"
            exit 1
            ;;
    esac
done

# Functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[OK]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

check_dependencies() {
    log_info "Checking dependencies..."

    local missing=()

    if ! command -v git &> /dev/null; then
        missing+=("git")
    fi

    if ! command -v tmux &> /dev/null; then
        missing+=("tmux")
    fi

    if [ ${#missing[@]} -ne 0 ]; then
        log_error "Missing dependencies: ${missing[*]}"
        log_info "Install with: apt-get install ${missing[*]}"
        exit 1
    fi

    log_success "All dependencies available"
}

check_git_repo() {
    log_info "Checking if $PROJECT_ROOT is a git repository..."

    if [ ! -d "$PROJECT_ROOT/.git" ]; then
        log_error "$PROJECT_ROOT is not a git repository"
        exit 1
    fi

    log_success "Git repository found"
}

check_base_branch() {
    log_info "Checking for base branch: $BASE_BRANCH..."

    cd "$PROJECT_ROOT"

    if ! git rev-parse --verify "$BASE_BRANCH" &> /dev/null; then
        log_error "Base branch '$BASE_BRANCH' does not exist"
        log_info "Available branches:"
        git branch -a
        exit 1
    fi

    log_success "Base branch '$BASE_BRANCH' found"
}

load_config() {
    if [ -n "$CONFIG_FILE" ] && [ -f "$CONFIG_FILE" ]; then
        log_info "Loading configuration from $CONFIG_FILE..."

        # Clear default agents
        unset AGENTS
        declare -gA AGENTS

        # Read config file (simple KEY=VALUE format)
        while IFS='=' read -r key value; do
            # Skip comments and empty lines
            [[ $key =~ ^#.*$ ]] && continue
            [[ -z $key ]] && continue

            AGENTS["$key"]="$value"
        done < "$CONFIG_FILE"

        log_success "Loaded ${#AGENTS[@]} agent configurations"
    fi
}

create_branch_if_needed() {
    local branch=$1

    cd "$PROJECT_ROOT"

    if git rev-parse --verify "$branch" &> /dev/null; then
        log_info "Branch '$branch' already exists"
        return 0
    fi

    log_info "Creating branch: $branch from $BASE_BRANCH"

    if [ "$DRY_RUN" = true ]; then
        echo "  Would run: git branch $branch $BASE_BRANCH"
    else
        git branch "$branch" "$BASE_BRANCH"
        log_success "Created branch: $branch"
    fi
}

create_worktree() {
    local agent_id=$1
    local branch=$2
    local worktree_path="$PROJECT_ROOT/../agent-$agent_id"

    cd "$PROJECT_ROOT"

    # Check if worktree already exists
    if [ -d "$worktree_path" ]; then
        log_warn "Worktree for '$agent_id' already exists at $worktree_path"
        return 0
    fi

    log_info "Creating worktree for agent: $agent_id"
    log_info "  Path: $worktree_path"
    log_info "  Branch: $branch"

    if [ "$DRY_RUN" = true ]; then
        echo "  Would run: git worktree add $worktree_path $branch"
    else
        git worktree add "$worktree_path" "$branch"
        log_success "Created worktree: $worktree_path"
    fi
}

create_tmux_session() {
    local agent_id=$1
    local worktree_path="$PROJECT_ROOT/../agent-$agent_id"

    # Check if session already exists
    if tmux has-session -t "$agent_id" 2>/dev/null; then
        log_warn "tmux session '$agent_id' already exists"
        return 0
    fi

    log_info "Creating tmux session: $agent_id"
    log_info "  Working directory: $worktree_path"

    if [ "$DRY_RUN" = true ]; then
        echo "  Would run: tmux new-session -d -s $agent_id -c $worktree_path"
    else
        tmux new-session -d -s "$agent_id" -c "$worktree_path"

        # Set up the session with a welcome message
        tmux send-keys -t "$agent_id" "echo 'Agent workspace: $agent_id'" Enter
        tmux send-keys -t "$agent_id" "echo 'Branch: $(cd "$worktree_path" && git branch --show-current)'" Enter
        tmux send-keys -t "$agent_id" "echo 'Ready for agent commands...'" Enter

        log_success "Created tmux session: $agent_id"
    fi
}

setup_agent() {
    local agent_id=$1
    local branch=$2

    echo ""
    log_info "Setting up agent: $agent_id"
    echo "============================================"

    # Create branch
    create_branch_if_needed "$branch"

    # Create worktree
    create_worktree "$agent_id" "$branch"

    # Create tmux session
    create_tmux_session "$agent_id"

    echo "============================================"
    log_success "Agent $agent_id setup complete"
}

list_agents() {
    echo ""
    log_info "Configured agents:"
    echo "============================================"
    for agent_id in "${!AGENTS[@]}"; do
        echo "  $agent_id -> ${AGENTS[$agent_id]}"
    done
    echo "============================================"
}

cleanup_agent() {
    local agent_id=$1
    local worktree_path="$PROJECT_ROOT/../agent-$agent_id"

    log_info "Cleaning up agent: $agent_id"

    # Kill tmux session
    if tmux has-session -t "$agent_id" 2>/dev/null; then
        if [ "$DRY_RUN" = true ]; then
            echo "  Would run: tmux kill-session -t $agent_id"
        else
            tmux kill-session -t "$agent_id"
            log_success "Killed tmux session: $agent_id"
        fi
    fi

    # Remove worktree
    if [ -d "$worktree_path" ]; then
        if [ "$DRY_RUN" = true ]; then
            echo "  Would run: git worktree remove $worktree_path"
        else
            cd "$PROJECT_ROOT"
            git worktree remove "$worktree_path" --force 2>/dev/null || true
            log_success "Removed worktree: $worktree_path"
        fi
    fi
}

show_status() {
    echo ""
    log_info "Current status:"
    echo "============================================"

    echo ""
    echo "Git Worktrees:"
    cd "$PROJECT_ROOT"
    git worktree list

    echo ""
    echo "tmux Sessions:"
    tmux list-sessions 2>/dev/null || echo "  No tmux sessions"

    echo "============================================"
}

# Main execution
main() {
    echo ""
    echo "========================================"
    echo "  Multi-Agent Workspace Setup"
    echo "========================================"
    echo ""

    if [ "$DRY_RUN" = true ]; then
        log_warn "DRY RUN MODE - No changes will be made"
        echo ""
    fi

    # Check dependencies
    check_dependencies

    # Check project
    check_git_repo
    check_base_branch

    # Load config if provided
    load_config

    # Show configured agents
    list_agents

    # Set up each agent
    for agent_id in "${!AGENTS[@]}"; do
        setup_agent "$agent_id" "${AGENTS[$agent_id]}"
    done

    # Show final status
    show_status

    echo ""
    log_success "All agents set up successfully!"
    echo ""
    echo "Quick commands:"
    echo "  - List sessions:    tmux list-sessions"
    echo "  - Attach to agent:  tmux attach -t <agent-id>"
    echo "  - Switch session:   tmux switch-client -t <agent-id>"
    echo "  - Kill session:     tmux kill-session -t <agent-id>"
    echo ""
}

# Run main
main
