#!/bin/bash
# claude-splitly-session-v2 - Worktree-based session manager for splitly

set -e

MAIN_REPO=~/projects/SafeSplit
WORKTREE_DIR=~/.worktrees/splitly
SESSION_DIR=~/.splitly-sessions
ENV_TEMPLATE=~/.env-templates/splitly
DEFAULT_BRANCH="main"
GITHUB_REPO="UniverseOfThings/SafeSplit"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

usage() {
    echo "Usage: ~/claude-splitly-session-v2 <command> [args]"
    echo ""
    echo "Commands:"
    echo "  new [ISSUE-NUM] [--open]  Create new session (optionally with GitHub issue)"
    echo "  ls                        List all sessions"
    echo "  attach <session>          Attach to session"
    echo "  kill <session>            Kill session and remove worktree"
    echo "  cleanup                   Kill sessions where PR is merged or issue closed"
    echo "  update                    Update current session metadata"
    exit 1
}

get_next_session_number() {
    local max=0
    for session in $(tmux list-sessions -F "#{session_name}" 2>/dev/null | grep "^splitly-[0-9]*$" | sed 's/splitly-//'); do
        if [ "$session" -gt "$max" ] 2>/dev/null; then
            max=$session
        fi
    done
    echo $((max + 1))
}

cmd_new() {
    local issue_num=""
    local open_warp=false
    
    for arg in "$@"; do
        if [ "$arg" = "--open" ]; then
            open_warp=true
        elif [ -z "$issue_num" ] && [ "$arg" != "--open" ]; then
            issue_num="$arg"
        fi
    done
    
    local session_num=$(get_next_session_number)
    local session_name="splitly-$session_num"
    local worktree_path="$WORKTREE_DIR/$session_name"
    
    echo -e "${BLUE}Creating session: $session_name${NC}"
    
    cd "$MAIN_REPO"
    git fetch origin --quiet
    
    if [ -n "$issue_num" ]; then
        local branch_name="feat/issue-$issue_num"
        git worktree add -b "$branch_name" "$worktree_path" "origin/$DEFAULT_BRANCH" 2>/dev/null || \
        git worktree add "$worktree_path" "origin/$DEFAULT_BRANCH"
    else
        git worktree add "$worktree_path" "origin/$DEFAULT_BRANCH" --detach
    fi
    
    echo -e "${GREEN}✓ Created worktree at $worktree_path${NC}"
    
    # Symlink env files
    [ -f "$ENV_TEMPLATE/.env" ] && ln -sf "$ENV_TEMPLATE/.env" "$worktree_path/.env"
    [ -f "$ENV_TEMPLATE/.envrc" ] && ln -sf "$ENV_TEMPLATE/.envrc" "$worktree_path/.envrc" && \
        (cd "$worktree_path" && direnv allow 2>/dev/null || true)

    # Symlink .claude directory for MCP access (remove if exists as dir)
    [ -d "$MAIN_REPO/.claude" ] && rm -rf "$worktree_path/.claude" && ln -s "$MAIN_REPO/.claude" "$worktree_path/.claude"

    # Add Rube MCP before starting session
    (cd "$worktree_path" && claude mcp add rube --transport http https://rube.app/mcp &>/dev/null) || true

    # Create tmux session and start Claude (unset CLAUDECODE to avoid nested session detection)
    tmux new-session -d -s "$session_name" -c "$worktree_path" -e "SPLITLY_SESSION=$session_name" -e "DIRENV_LOG_FORMAT="
    tmux send-keys -t "$session_name" "unset CLAUDECODE && claude --dangerously-skip-permissions" Enter
    
    # Create metadata file
    mkdir -p "$SESSION_DIR"
    cat > "$SESSION_DIR/$session_name" << EOF
worktree=$worktree_path
branch=$(cd "$worktree_path" && git branch --show-current || echo "detached")
status=starting
EOF
    [ -n "$issue_num" ] && echo "issue=$issue_num" >> "$SESSION_DIR/$session_name"
    
    echo -e "${GREEN}✓ Created tmux session: $session_name${NC}"
    echo -e "${GREEN}✓ Started Claude in session${NC}"
    echo ""
    echo "Attach with: tmux attach -t $session_name"
    echo "Or open in iTerm2: ~/open-iterm-tab $session_name"

    # Clean output for scripting
    echo "SESSION=$session_name"

    if [ "$open_warp" = true ]; then
        ~/open-iterm-tab "$session_name"
    fi
}

cmd_ls() {
    echo -e "${BLUE}Splitly Sessions:${NC}"
    echo ""
    
    for session in $(tmux list-sessions -F "#{session_name}" 2>/dev/null | grep "^splitly-" | sort -V); do
        local meta_file="$SESSION_DIR/$session"
        local worktree="" branch="" status="" pr="" issue=""
        
        if [ -f "$meta_file" ]; then
            worktree=$(grep "^worktree=" "$meta_file" 2>/dev/null | cut -d= -f2-)
            branch=$(grep "^branch=" "$meta_file" 2>/dev/null | cut -d= -f2-)
            status=$(grep "^status=" "$meta_file" 2>/dev/null | cut -d= -f2-)
            pr=$(grep "^pr=" "$meta_file" 2>/dev/null | cut -d= -f2-)
            issue=$(grep "^issue=" "$meta_file" 2>/dev/null | cut -d= -f2-)
        fi
        
        [ -d "$worktree" ] && branch=$(cd "$worktree" && git branch --show-current 2>/dev/null || echo "$branch")
        
        echo -e "${GREEN}$session${NC}"
        [ -n "$branch" ] && echo "  Branch: $branch"
        [ -n "$issue" ] && echo "  Issue: #$issue"
        [ -n "$status" ] && echo "  Status: $status"
        [ -n "$pr" ] && echo "  PR: $pr"
        echo ""
    done
}

cmd_attach() {
    [ -z "$1" ] && echo "Usage: ~/claude-splitly-session-v2 attach <session>" && exit 1
    tmux attach -t "$1"
}

cmd_kill() {
    local session="$1"
    [ -z "$session" ] && echo "Usage: ~/claude-splitly-session-v2 kill <session>" && exit 1
    
    local meta_file="$SESSION_DIR/$session"
    local worktree=""
    [ -f "$meta_file" ] && worktree=$(grep "^worktree=" "$meta_file" 2>/dev/null | cut -d= -f2-)
    
    tmux kill-session -t "$session" 2>/dev/null && echo -e "${GREEN}✓ Killed tmux session: $session${NC}"
    
    if [ -n "$worktree" ] && [ -d "$worktree" ]; then
        cd "$MAIN_REPO"
        git worktree remove --force "$worktree" 2>/dev/null && echo -e "${GREEN}✓ Removed worktree: $worktree${NC}"
    fi
    
    # Archive metadata instead of deleting
    if [ -f "$meta_file" ]; then
        mkdir -p "$SESSION_DIR/archive"
        mv "$meta_file" "$SESSION_DIR/archive/${session}_$(date +%Y%m%d-%H%M%S)"
        echo -e "${GREEN}✓ Archived metadata${NC}"
    fi
}

cmd_cleanup() {
    echo "Checking for completed sessions..."
    
    for session in $(tmux list-sessions -F "#{session_name}" 2>/dev/null | grep "^splitly-"); do
        local meta_file="$SESSION_DIR/$session"
        local pr="" issue=""
        
        [ -f "$meta_file" ] && {
            pr=$(grep "^pr=" "$meta_file" 2>/dev/null | cut -d= -f2-)
            issue=$(grep "^issue=" "$meta_file" 2>/dev/null | cut -d= -f2-)
        }
        
        if [ -n "$pr" ]; then
            local pr_num=$(echo "$pr" | grep -o '[0-9]*$')
            if [ -n "$pr_num" ]; then
                local state=$(gh pr view "$pr_num" --repo "$GITHUB_REPO" --json state -q '.state' 2>/dev/null)
                [ "$state" = "MERGED" ] && echo -e "${YELLOW}PR #$pr_num merged - killing $session${NC}" && cmd_kill "$session" && continue
            fi
        fi
        
        if [ -n "$issue" ]; then
            local issue_state=$(gh issue view "$issue" --repo "$GITHUB_REPO" --json state -q '.state' 2>/dev/null)
            [ "$issue_state" = "CLOSED" ] && echo -e "${YELLOW}Issue #$issue closed - killing $session${NC}" && cmd_kill "$session"
        fi
    done
    echo "Cleanup complete."
}

cmd_update() {
    local session="${SPLITLY_SESSION:-}"
    [ -z "$session" ] && echo "Error: Not in a splitly session" && exit 1
    
    local meta_file="$SESSION_DIR/$session"
    local worktree=$(grep "^worktree=" "$meta_file" 2>/dev/null | cut -d= -f2-)
    [ -z "$worktree" ] || [ ! -d "$worktree" ] && worktree="$PWD"
    
    local branch=$(cd "$worktree" && git branch --show-current 2>/dev/null)
    [ -n "$branch" ] && sed -i '' "s|^branch=.*|branch=$branch|" "$meta_file" 2>/dev/null
    
    echo "Updated metadata for $session"
}

case "${1:-}" in
    new)     shift; cmd_new "$@" ;;
    ls)      cmd_ls ;;
    attach)  cmd_attach "$2" ;;
    kill)    cmd_kill "$2" ;;
    cleanup) cmd_cleanup ;;
    update)  cmd_update ;;
    *)       usage ;;
esac