#!/bin/bash
###############################################################################
# Setup Worktree with Port Allocation
#
# Creates a git worktree with automatically allocated ports for parallel
# development. Each worktree gets its own PORT and CHROME_WS_PORT to avoid
# conflicts when running multiple instances simultaneously.
#
# Usage: ./scripts/setup-worktree.sh <worktree-name> <branch-name>
# Example: ./scripts/setup-worktree.sh my-feature feature/my-feature
#
# This script:
# 1. Auto-cleans stale port allocations from removed worktrees
# 2. Fetches latest from origin to ensure branch availability
# 3. Creates the git worktree in .worktrees/<name>
# 4. Allocates the next available port from .worktree-ports.json
# 5. Creates .env.local in the worktree with proper PORT settings
# 6. Copies other env vars from main repo's .env.local
###############################################################################

set -e

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

# Get script directory and project root
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
PORTS_FILE="$PROJECT_ROOT/.worktree-ports.json"
WORKTREES_DIR="$PROJECT_ROOT/.worktrees"

# Validate arguments
if [ $# -lt 2 ]; then
  echo -e "${RED}Usage: $0 <worktree-name> <branch-name>${NC}"
  echo ""
  echo "Examples:"
  echo "  $0 my-feature feature/my-feature     # Create new branch"
  echo "  $0 bugfix-123 fix/issue-123          # Create new branch"
  echo "  $0 review-pr main                    # Use existing branch"
  exit 1
fi

WORKTREE_NAME="$1"
BRANCH_NAME="$2"
WORKTREE_PATH="$WORKTREES_DIR/$WORKTREE_NAME"

###############################################################################
# Cleanup stale port allocations
#
# Removes entries from .worktree-ports.json for worktrees that no longer exist,
# then recalculates the next available port to prevent port number inflation.
###############################################################################
cleanup_stale_allocations() {
  if [ ! -f "$PORTS_FILE" ]; then
    return 0
  fi

  echo "Checking for stale port allocations..."

  # Use node to clean up and recalculate
  local cleaned=$(node -e "
const fs = require('fs');
const ports = require('$PORTS_FILE');

let removed = 0;
let maxPort = 3000; // Base port

// Check each allocation
const newAllocated = {};
for (const [path, config] of Object.entries(ports.allocated)) {
  if (fs.existsSync(path)) {
    newAllocated[path] = config;
    if (config.appPort > maxPort) {
      maxPort = config.appPort;
    }
  } else {
    removed++;
    console.error('  Removed stale: ' + path.split('/').pop() + ' (port ' + config.appPort + ')');
  }
}

if (removed > 0) {
  ports.allocated = newAllocated;
  ports.nextAppPort = maxPort + 1;
  ports.nextChromePort = 9222 + (maxPort - 3000) + 1;
  fs.writeFileSync('$PORTS_FILE', JSON.stringify(ports, null, 2) + '\n');
}

console.log(removed);
" 2>&1)

  # Parse output - last line is the count, earlier lines are removed entries
  local count=$(echo "$cleaned" | tail -1)
  local messages=$(echo "$cleaned" | head -n -1)

  if [ "$count" -gt 0 ] 2>/dev/null; then
    echo -e "${YELLOW}Cleaned up $count stale allocation(s):${NC}"
    echo "$messages"
    echo ""
  fi
}

# Run cleanup before creating new worktree
cleanup_stale_allocations

# Check if worktree already exists
if [ -d "$WORKTREE_PATH" ]; then
  echo -e "${RED}Error: Worktree '$WORKTREE_NAME' already exists at $WORKTREE_PATH${NC}"
  exit 1
fi

# Initialize .worktree-ports.json if it doesn't exist
if [ ! -f "$PORTS_FILE" ]; then
  echo -e "${YELLOW}Initializing .worktree-ports.json...${NC}"
  cat > "$PORTS_FILE" << 'EOF'
{
  "description": "Port allocation for parallel worktree development. Chrome debug port = 9222 + (appPort - 3000)",
  "baseAppPort": 3000,
  "baseChromePort": 9222,
  "allocated": {},
  "nextAppPort": 3001,
  "nextChromePort": 9223
}
EOF
  # Add main repo allocation
  MAIN_ENTRY="\"$PROJECT_ROOT\": { \"appPort\": 3000, \"chromePort\": 9222 }"
  # Use a temp file for sed compatibility
  sed -i "s|\"allocated\": {}|\"allocated\": { $MAIN_ENTRY }|" "$PORTS_FILE"
  echo -e "${GREEN}Created .worktree-ports.json${NC}"
fi

# Read next available ports using node (for JSON parsing)
NEXT_APP_PORT=$(node -e "console.log(require('$PORTS_FILE').nextAppPort)")
NEXT_CHROME_PORT=$(node -e "console.log(require('$PORTS_FILE').nextChromePort)")

echo ""
echo -e "${GREEN}Setting up worktree: $WORKTREE_NAME${NC}"
echo "  Branch: $BRANCH_NAME"
echo "  Path: $WORKTREE_PATH"
echo "  App Port: $NEXT_APP_PORT"
echo "  Chrome Debug Port: $NEXT_CHROME_PORT"
echo ""

# Create worktrees directory if needed
mkdir -p "$WORKTREES_DIR"

# Fetch latest from origin
echo "Fetching latest from origin..."
git fetch origin --quiet

# Check if branch exists (local or remote)
if git rev-parse --verify "$BRANCH_NAME" >/dev/null 2>&1; then
  echo "Using existing branch: $BRANCH_NAME"
  git worktree add "$WORKTREE_PATH" "$BRANCH_NAME"
elif git rev-parse --verify "origin/$BRANCH_NAME" >/dev/null 2>&1; then
  echo "Using remote branch: origin/$BRANCH_NAME"
  git worktree add --track -b "$BRANCH_NAME" "$WORKTREE_PATH" "origin/$BRANCH_NAME"
else
  echo "Creating new branch: $BRANCH_NAME (from origin/dev)"
  git worktree add -b "$BRANCH_NAME" "$WORKTREE_PATH" origin/dev
fi

# Create .env.local in the worktree
echo ""
echo "Creating .env.local in worktree..."

# Start with port configuration
cat > "$WORKTREE_PATH/.env.local" << EOF
# Worktree-specific port configuration
# Generated by setup-worktree.sh
PORT=$NEXT_APP_PORT
NEXT_PUBLIC_BASE_URL=http://localhost:$NEXT_APP_PORT
NEXTAUTH_URL=http://localhost:$NEXT_APP_PORT
CHROME_WS_PORT=$NEXT_CHROME_PORT

EOF

# Copy other env vars from main repo's .env.local (excluding PORT-related ones)
if [ -f "$PROJECT_ROOT/.env.local" ]; then
  echo "Copying environment variables from main repo..."
  grep -v -E "^(PORT|NEXT_PUBLIC_BASE_URL|NEXTAUTH_URL|CHROME_WS_PORT)=" "$PROJECT_ROOT/.env.local" >> "$WORKTREE_PATH/.env.local" 2>/dev/null || true
  echo -e "${GREEN}Copied env vars from main .env.local${NC}"
else
  echo -e "${YELLOW}Warning: No .env.local found in main repo. You'll need to add DATABASE_URL, STRIPE keys, etc. manually.${NC}"
fi

# Update .worktree-ports.json with new allocation
echo ""
echo "Updating port registry..."

# Use node to update the JSON file properly
node -e "
const fs = require('fs');
const ports = require('$PORTS_FILE');
ports.allocated['$WORKTREE_PATH'] = { appPort: $NEXT_APP_PORT, chromePort: $NEXT_CHROME_PORT };
ports.nextAppPort = $NEXT_APP_PORT + 1;
ports.nextChromePort = $NEXT_CHROME_PORT + 1;
fs.writeFileSync('$PORTS_FILE', JSON.stringify(ports, null, 2) + '\n');
"

echo -e "${GREEN}Port registry updated${NC}"

# Final instructions
echo ""
echo "==========================================================================="
echo -e "${GREEN}Worktree setup complete!${NC}"
echo "==========================================================================="
echo ""
echo "To start working:"
echo "  cd $WORKTREE_PATH"
echo "  npm run power-cycle    # Start dev server on port $NEXT_APP_PORT"
echo ""
echo "To start Claude Code with browser automation:"
echo "  cd $WORKTREE_PATH"
echo "  CHROME_WS_PORT=$NEXT_CHROME_PORT claude"
echo ""
echo "Your app will be available at: http://localhost:$NEXT_APP_PORT"
echo "==========================================================================="